Generate Assertion Ops for interpolate_bilinear

Author: faizan-m

tensorflow_addons/image/dense_image_warp.py

@@ -52,27 +52,22 @@ def interpolate_bilinear(
     with tf.name_scope(name or "interpolate_bilinear"):
         grid = tf.convert_to_tensor(grid)
         query_points = tf.convert_to_tensor(query_points)
-
-        # grid shape checks
-        grid_static_shape = grid.shape
         grid_shape = tf.shape(grid)
-        if grid_static_shape.dims is not None:
-            if len(grid_static_shape) != 4:
-                raise ValueError("Grid must be 4D Tensor")
-            if grid_static_shape[1] is not None and grid_static_shape[1] < 2:
-                raise ValueError("Grid height must be at least 2.")
-            if grid_static_shape[2] is not None and grid_static_shape[2] < 2:
-                raise ValueError("Grid width must be at least 2.")
-
-        # query_points shape checks
-        query_static_shape = query_points.shape
         query_shape = tf.shape(query_points)
-        if query_static_shape.dims is not None:
-            if len(query_static_shape) != 3:
-                raise ValueError("Query points must be 3 dimensional.")
-            query_hw = query_static_shape[2]
-            if query_hw is not None and query_hw != 2:
-                raise ValueError("Query points last dimension must be 2.")
+
+        tf.Assert(tf.equal(tf.rank(grid), 4), ["Grid must be 4D Tensor"])
+        tf.Assert(
+            tf.greater_equal(grid_shape[1], 2), ["Grid height must be at least 2."]
+        )
+        tf.Assert(
+            tf.greater_equal(grid_shape[2], 2), ["Grid width must be at least 2."]
+        )
+        tf.Assert(
+            tf.equal(tf.rank(query_points), 3), ["Query points must be 3 dimensional."]
+        )
+        tf.Assert(
+            tf.equal(query_shape[2], 2), ["Query points last dimension must be 2."]
+        )
 
         batch_size, height, width, channels = (
             grid_shape[0],

tensorflow_addons/image/tests/dense_image_warp_test.py

@@ -239,7 +239,9 @@ def test_interpolation():
 def test_size_exception():
     """Make sure it throws an exception for images that are too small."""
     shape = [1, 2, 1, 1]
-    with pytest.raises(ValueError, match="Grid width must be at least 2."):
+    with pytest.raises(
+        tf.errors.InvalidArgumentError, match="Grid width must be at least 2."
+    ):
         _check_interpolation_correctness(shape, "float32", "float32")
 
 
@@ -250,11 +252,3 @@ def test_unknown_shapes():
     shapes_to_try = [[3, 4, 5, 6], [1, 2, 2, 1]]
     for shape in shapes_to_try:
         _check_interpolation_correctness(shape, "float32", "float32", True)
-
-
-@pytest.mark.usefixtures("only_run_functions_eagerly")
-def test_symbolic_tensor_shape():
-    image = tf.keras.layers.Input(shape=(7, 7, 192))
-    flow = tf.ones((1, 7, 7, 2))
-    interp = dense_image_warp(image, flow)
-    np.testing.assert_array_equal(interp.shape.as_list(), [None, 7, 7, 192])