We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
(None, None, None, None)-shaped variable returned by tfa.image.dense_image_warp
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux 20.04
- TensorFlow version and how it was installed (source or binary): 2.1.0, source
- TensorFlow-Addons version and how it was installed (source or binary): latest, source
- Python version: 3.7
- Is GPU used? (yes/no): Yes

**Describe the bug**

When passing a "None"-batch size image with flow, such as a (None, 7, 7, 192) image with a valid flow, the returned value will have shape (None, None, None, None) 

**Code to reproduce the issue**
```
import tensorflow as tf
import tensorflow_addons as tfa

tf.config.experimental_run_functions_eagerly(True)

image = tf.keras.layers.Input(shape=(7, 7, 192))
flow = tf.ones((1, 7, 7, 2))
warped = tfa.image.dense_image_warp(image, flow)
```

**Other info / logs**
```
>>> warped.shape
TensorShape([None, None, None, None])
```

Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/conftest.py`
Content:
```
1 from tensorflow_addons.utils.test_utils import (  # noqa: F401
2     maybe_run_functions_eagerly,
3     run_custom_and_py_ops,
4     run_with_mixed_precision_policy,
5     pytest_make_parametrize_id,
6     data_format,
7     set_seeds,
8     pytest_addoption,
9     set_global_variables,
10     pytest_configure,
11     device,
12     pytest_generate_tests,
13     pytest_collection_modifyitems,
14 )
15 
16 # fixtures present in this file will be available
17 # when running tests and can be referenced with strings
18 # https://docs.pytest.org/en/latest/fixture.html#conftest-py-sharing-fixture-functions
19
```

Path: `tensorflow_addons/image/dense_image_warp.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 """Image warping using per-pixel flow vectors."""
16 
17 import numpy as np
18 import tensorflow as tf
19 
20 from tensorflow_addons.utils import types
21 from typing import Optional
22 
23 
24 @tf.function
25 def interpolate_bilinear(
26     grid: types.TensorLike,
27     query_points: types.TensorLike,
28     indexing: str = "ij",
29     name: Optional[str] = None,
30 ) -> tf.Tensor:
31     """Similar to Matlab's interp2 function.
32 
33     Finds values for query points on a grid using bilinear interpolation.
34 
35     Args:
36       grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.
37       query_points: a 3-D float `Tensor` of N points with shape
38         `[batch, N, 2]`.
39       indexing: whether the query points are specified as row and column (ij),
40         or Cartesian coordinates (xy).
41       name: a name for the operation (optional).
42 
43     Returns:
44       values: a 3-D `Tensor` with shape `[batch, N, channels]`
45 
46     Raises:
47       ValueError: if the indexing mode is invalid, or if the shape of the
48         inputs invalid.
49     """
50     if indexing != "ij" and indexing != "xy":
51         raise ValueError("Indexing mode must be 'ij' or 'xy'")
52 
53     with tf.name_scope(name or "interpolate_bilinear"):
54         grid = tf.convert_to_tensor(grid)
55         query_points = tf.convert_to_tensor(query_points)
56 
57         # grid shape checks
58         grid_static_shape = grid.shape
59         grid_shape = tf.shape(grid)
60         if grid_static_shape.dims is not None:
61             if len(grid_static_shape) != 4:
62                 raise ValueError("Grid must be 4D Tensor")
63             if grid_static_shape[1] is not None and grid_static_shape[1] < 2:
64                 raise ValueError("Grid height must be at least 2.")
65             if grid_static_shape[2] is not None and grid_static_shape[2] < 2:
66                 raise ValueError("Grid width must be at least 2.")
67         else:
68             with tf.control_dependencies(
69                 [
70                     tf.debugging.assert_greater_equal(
71                         grid_shape[1], 2, message="Grid height must be at least 2."
72                     ),
73                     tf.debugging.assert_greater_equal(
74                         grid_shape[2], 2, message="Grid width must be at least 2."
75                     ),
76                     tf.debugging.assert_less_equal(
77                         tf.cast(
78                             grid_shape[0] * grid_shape[1] * grid_shape[2],
79                             dtype=tf.dtypes.float32,
80                         ),
81                         np.iinfo(np.int32).max / 8.0,
82                         message="The image size or batch size is sufficiently "
83                         "large that the linearized addresses used by "
84                         "tf.gather may exceed the int32 limit.",
85                     ),
86                 ]
87             ):
88                 pass
89 
90         # query_points shape checks
91         query_static_shape = query_points.shape
92         query_shape = tf.shape(query_points)
93         if query_static_shape.dims is not None:
94             if len(query_static_shape) != 3:
95                 raise ValueError("Query points must be 3 dimensional.")
96             query_hw = query_static_shape[2]
97             if query_hw is not None and query_hw != 2:
98                 raise ValueError("Query points last dimension must be 2.")
99         else:
100             with tf.control_dependencies(
101                 [
102                     tf.debugging.assert_equal(
103                         query_shape[2],
104                         2,
105                         message="Query points last dimension must be 2.",
106                     )
107                 ]
108             ):
109                 pass
110 
111         batch_size, height, width, channels = (
112             grid_shape[0],
113             grid_shape[1],
114             grid_shape[2],
115             grid_shape[3],
116         )
117 
118         num_queries = query_shape[1]
119 
120         query_type = query_points.dtype
121         grid_type = grid.dtype
122 
123         alphas = []
124         floors = []
125         ceils = []
126         index_order = [0, 1] if indexing == "ij" else [1, 0]
127         unstacked_query_points = tf.unstack(query_points, axis=2, num=2)
128 
129         for i, dim in enumerate(index_order):
130             with tf.name_scope("dim-" + str(dim)):
131                 queries = unstacked_query_points[dim]
132 
133                 size_in_indexing_dimension = grid_shape[i + 1]
134 
135                 # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1
136                 # is still a valid index into the grid.
137                 max_floor = tf.cast(size_in_indexing_dimension - 2, query_type)
138                 min_floor = tf.constant(0.0, dtype=query_type)
139                 floor = tf.math.minimum(
140                     tf.math.maximum(min_floor, tf.math.floor(queries)), max_floor
141                 )
142                 int_floor = tf.cast(floor, tf.dtypes.int32)
143                 floors.append(int_floor)
144                 ceil = int_floor + 1
145                 ceils.append(ceil)
146 
147                 # alpha has the same type as the grid, as we will directly use alpha
148                 # when taking linear combinations of pixel values from the image.
149                 alpha = tf.cast(queries - floor, grid_type)
150                 min_alpha = tf.constant(0.0, dtype=grid_type)
151                 max_alpha = tf.constant(1.0, dtype=grid_type)
152                 alpha = tf.math.minimum(tf.math.maximum(min_alpha, alpha), max_alpha)
153 
154                 # Expand alpha to [b, n, 1] so we can use broadcasting
155                 # (since the alpha values don't depend on the channel).
156                 alpha = tf.expand_dims(alpha, 2)
157                 alphas.append(alpha)
158 
159             flattened_grid = tf.reshape(grid, [batch_size * height * width, channels])
160             batch_offsets = tf.reshape(
161                 tf.range(batch_size) * height * width, [batch_size, 1]
162             )
163 
164         # This wraps tf.gather. We reshape the image data such that the
165         # batch, y, and x coordinates are pulled into the first dimension.
166         # Then we gather. Finally, we reshape the output back. It's possible this
167         # code would be made simpler by using tf.gather_nd.
168         def gather(y_coords, x_coords, name):
169             with tf.name_scope("gather-" + name):
170                 linear_coordinates = batch_offsets + y_coords * width + x_coords
171                 gathered_values = tf.gather(flattened_grid, linear_coordinates)
172                 return tf.reshape(gathered_values, [batch_size, num_queries, channels])
173 
174         # grab the pixel values in the 4 corners around each query point
175         top_left = gather(floors[0], floors[1], "top_left")
176         top_right = gather(floors[0], ceils[1], "top_right")
177         bottom_left = gather(ceils[0], floors[1], "bottom_left")
178         bottom_right = gather(ceils[0], ceils[1], "bottom_right")
179 
180         # now, do the actual interpolation
181         with tf.name_scope("interpolate"):
182             interp_top = alphas[1] * (top_right - top_left) + top_left
183             interp_bottom = alphas[1] * (bottom_right - bottom_left) + bottom_left
184             interp = alphas[0] * (interp_bottom - interp_top) + interp_top
185 
186         return interp
187 
188 
189 @tf.function
190 def dense_image_warp(
191     image: types.TensorLike, flow: types.TensorLike, name: Optional[str] = None
192 ) -> tf.Tensor:
193     """Image warping using per-pixel flow vectors.
194 
195     Apply a non-linear warp to the image, where the warp is specified by a
196     dense flow field of offset vectors that define the correspondences of
197     pixel values in the output image back to locations in the source image.
198     Specifically, the pixel value at `output[b, j, i, c]` is
199     `images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c]`.
200 
201     The locations specified by this formula do not necessarily map to an int
202     index. Therefore, the pixel value is obtained by bilinear
203     interpolation of the 4 nearest pixels around
204     `(b, j - flow[b, j, i, 0], i - flow[b, j, i, 1])`. For locations outside
205     of the image, we use the nearest pixel values at the image boundary.
206 
207     PLEASE NOTE: The definition of the flow field above is different from that
208     of optical flow. This function expects the negative forward flow from
209     output image to source image. Given two images `I_1` and `I_2` and the
210     optical flow `F_12` from `I_1` to `I_2`, the image `I_1` can be
211     reconstructed by `I_1_rec = dense_image_warp(I_2, -F_12)`.
212 
213     Args:
214       image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
215       flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
216       name: A name for the operation (optional).
217 
218       Note that image and flow can be of type tf.half, tf.float32, or
219       tf.float64, and do not necessarily have to be the same type.
220 
221     Returns:
222       A 4-D float `Tensor` with shape`[batch, height, width, channels]`
223         and same type as input image.
224 
225     Raises:
226       ValueError: if height < 2 or width < 2 or the inputs have the wrong
227         number of dimensions.
228     """
229     with tf.name_scope(name or "dense_image_warp"):
230         image = tf.convert_to_tensor(image)
231         flow = tf.convert_to_tensor(flow)
232         batch_size, height, width, channels = (
233             tf.shape(image)[0],
234             tf.shape(image)[1],
235             tf.shape(image)[2],
236             tf.shape(image)[3],
237         )
238 
239         # The flow is defined on the image grid. Turn the flow into a list of query
240         # points in the grid space.
241         grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))
242         stacked_grid = tf.cast(tf.stack([grid_y, grid_x], axis=2), flow.dtype)
243         batched_grid = tf.expand_dims(stacked_grid, axis=0)
244         query_points_on_grid = batched_grid - flow
245         query_points_flattened = tf.reshape(
246             query_points_on_grid, [batch_size, height * width, 2]
247         )
248         # Compute values at the query points, then reshape the result back to the
249         # image grid.
250         interpolated = interpolate_bilinear(image, query_points_flattened)
251         interpolated = tf.reshape(interpolated, [batch_size, height, width, channels])
252         return interpolated
253
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/tensorflow_addons/conftest.py b/tensorflow_addons/conftest.py
--- a/tensorflow_addons/conftest.py
+++ b/tensorflow_addons/conftest.py
@@ -1,5 +1,6 @@
 from tensorflow_addons.utils.test_utils import (  # noqa: F401
     maybe_run_functions_eagerly,
+    only_run_functions_eagerly,
     run_custom_and_py_ops,
     run_with_mixed_precision_policy,
     pytest_make_parametrize_id,
diff --git a/tensorflow_addons/image/dense_image_warp.py b/tensorflow_addons/image/dense_image_warp.py
--- a/tensorflow_addons/image/dense_image_warp.py
+++ b/tensorflow_addons/image/dense_image_warp.py
@@ -186,6 +186,12 @@
         return interp
 
 
+def _get_dim(x, idx):
+    if x.shape.ndims is None:
+        return tf.shape(x)[idx]
+    return x.shape[idx] or tf.shape(x)[idx]
+
+
 @tf.function
 def dense_image_warp(
     image: types.TensorLike, flow: types.TensorLike, name: Optional[str] = None
@@ -230,10 +236,10 @@
         image = tf.convert_to_tensor(image)
         flow = tf.convert_to_tensor(flow)
         batch_size, height, width, channels = (
-            tf.shape(image)[0],
-            tf.shape(image)[1],
-            tf.shape(image)[2],
-            tf.shape(image)[3],
+            _get_dim(image, 0),
+            _get_dim(image, 1),
+            _get_dim(image, 2),
+            _get_dim(image, 3),
         )
 
         # The flow is defined on the image grid. Turn the flow into a list of query

{"golden_diff": "diff --git a/tensorflow_addons/conftest.py b/tensorflow_addons/conftest.py\n--- a/tensorflow_addons/conftest.py\n+++ b/tensorflow_addons/conftest.py\n@@ -1,5 +1,6 @@\n from tensorflow_addons.utils.test_utils import (  # noqa: F401\n     maybe_run_functions_eagerly,\n+    only_run_functions_eagerly,\n     run_custom_and_py_ops,\n     run_with_mixed_precision_policy,\n     pytest_make_parametrize_id,\ndiff --git a/tensorflow_addons/image/dense_image_warp.py b/tensorflow_addons/image/dense_image_warp.py\n--- a/tensorflow_addons/image/dense_image_warp.py\n+++ b/tensorflow_addons/image/dense_image_warp.py\n@@ -186,6 +186,12 @@\n         return interp\n \n \n+def _get_dim(x, idx):\n+    if x.shape.ndims is None:\n+        return tf.shape(x)[idx]\n+    return x.shape[idx] or tf.shape(x)[idx]\n+\n+\n @tf.function\n def dense_image_warp(\n     image: types.TensorLike, flow: types.TensorLike, name: Optional[str] = None\n@@ -230,10 +236,10 @@\n         image = tf.convert_to_tensor(image)\n         flow = tf.convert_to_tensor(flow)\n         batch_size, height, width, channels = (\n-            tf.shape(image)[0],\n-            tf.shape(image)[1],\n-            tf.shape(image)[2],\n-            tf.shape(image)[3],\n+            _get_dim(image, 0),\n+            _get_dim(image, 1),\n+            _get_dim(image, 2),\n+            _get_dim(image, 3),\n         )\n \n         # The flow is defined on the image grid. Turn the flow into a list of query\n", "issue": "(None, None, None, None)-shaped variable returned by tfa.image.dense_image_warp\n**System information**\r\n- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux 20.04\r\n- TensorFlow version and how it was installed (source or binary): 2.1.0, source\r\n- TensorFlow-Addons version and how it was installed (source or binary): latest, source\r\n- Python version: 3.7\r\n- Is GPU used? (yes/no): Yes\r\n\r\n**Describe the bug**\r\n\r\nWhen passing a \"None\"-batch size image with flow, such as a (None, 7, 7, 192) image with a valid flow, the returned value will have shape (None, None, None, None) \r\n\r\n**Code to reproduce the issue**\r\n```\r\nimport tensorflow as tf\r\nimport tensorflow_addons as tfa\r\n\r\ntf.config.experimental_run_functions_eagerly(True)\r\n\r\nimage = tf.keras.layers.Input(shape=(7, 7, 192))\r\nflow = tf.ones((1, 7, 7, 2))\r\nwarped = tfa.image.dense_image_warp(image, flow)\r\n```\r\n\r\n**Other info / logs**\r\n```\r\n>>> warped.shape\r\nTensorShape([None, None, None, None])\r\n```\r\n\r\nInclude any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.\r\n\n", "before_files": [{"content": "from tensorflow_addons.utils.test_utils import (  # noqa: F401\n    maybe_run_functions_eagerly,\n    run_custom_and_py_ops,\n    run_with_mixed_precision_policy,\n    pytest_make_parametrize_id,\n    data_format,\n    set_seeds,\n    pytest_addoption,\n    set_global_variables,\n    pytest_configure,\n    device,\n    pytest_generate_tests,\n    pytest_collection_modifyitems,\n)\n\n# fixtures present in this file will be available\n# when running tests and can be referenced with strings\n# https://docs.pytest.org/en/latest/fixture.html#conftest-py-sharing-fixture-functions\n", "path": "tensorflow_addons/conftest.py"}, {"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Image warping using per-pixel flow vectors.\"\"\"\n\nimport numpy as np\nimport tensorflow as tf\n\nfrom tensorflow_addons.utils import types\nfrom typing import Optional\n\n\n@tf.function\ndef interpolate_bilinear(\n    grid: types.TensorLike,\n    query_points: types.TensorLike,\n    indexing: str = \"ij\",\n    name: Optional[str] = None,\n) -> tf.Tensor:\n    \"\"\"Similar to Matlab's interp2 function.\n\n    Finds values for query points on a grid using bilinear interpolation.\n\n    Args:\n      grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.\n      query_points: a 3-D float `Tensor` of N points with shape\n        `[batch, N, 2]`.\n      indexing: whether the query points are specified as row and column (ij),\n        or Cartesian coordinates (xy).\n      name: a name for the operation (optional).\n\n    Returns:\n      values: a 3-D `Tensor` with shape `[batch, N, channels]`\n\n    Raises:\n      ValueError: if the indexing mode is invalid, or if the shape of the\n        inputs invalid.\n    \"\"\"\n    if indexing != \"ij\" and indexing != \"xy\":\n        raise ValueError(\"Indexing mode must be 'ij' or 'xy'\")\n\n    with tf.name_scope(name or \"interpolate_bilinear\"):\n        grid = tf.convert_to_tensor(grid)\n        query_points = tf.convert_to_tensor(query_points)\n\n        # grid shape checks\n        grid_static_shape = grid.shape\n        grid_shape = tf.shape(grid)\n        if grid_static_shape.dims is not None:\n            if len(grid_static_shape) != 4:\n                raise ValueError(\"Grid must be 4D Tensor\")\n            if grid_static_shape[1] is not None and grid_static_shape[1] < 2:\n                raise ValueError(\"Grid height must be at least 2.\")\n            if grid_static_shape[2] is not None and grid_static_shape[2] < 2:\n                raise ValueError(\"Grid width must be at least 2.\")\n        else:\n            with tf.control_dependencies(\n                [\n                    tf.debugging.assert_greater_equal(\n                        grid_shape[1], 2, message=\"Grid height must be at least 2.\"\n                    ),\n                    tf.debugging.assert_greater_equal(\n                        grid_shape[2], 2, message=\"Grid width must be at least 2.\"\n                    ),\n                    tf.debugging.assert_less_equal(\n                        tf.cast(\n                            grid_shape[0] * grid_shape[1] * grid_shape[2],\n                            dtype=tf.dtypes.float32,\n                        ),\n                        np.iinfo(np.int32).max / 8.0,\n                        message=\"The image size or batch size is sufficiently \"\n                        \"large that the linearized addresses used by \"\n                        \"tf.gather may exceed the int32 limit.\",\n                    ),\n                ]\n            ):\n                pass\n\n        # query_points shape checks\n        query_static_shape = query_points.shape\n        query_shape = tf.shape(query_points)\n        if query_static_shape.dims is not None:\n            if len(query_static_shape) != 3:\n                raise ValueError(\"Query points must be 3 dimensional.\")\n            query_hw = query_static_shape[2]\n            if query_hw is not None and query_hw != 2:\n                raise ValueError(\"Query points last dimension must be 2.\")\n        else:\n            with tf.control_dependencies(\n                [\n                    tf.debugging.assert_equal(\n                        query_shape[2],\n                        2,\n                        message=\"Query points last dimension must be 2.\",\n                    )\n                ]\n            ):\n                pass\n\n        batch_size, height, width, channels = (\n            grid_shape[0],\n            grid_shape[1],\n            grid_shape[2],\n            grid_shape[3],\n        )\n\n        num_queries = query_shape[1]\n\n        query_type = query_points.dtype\n        grid_type = grid.dtype\n\n        alphas = []\n        floors = []\n        ceils = []\n        index_order = [0, 1] if indexing == \"ij\" else [1, 0]\n        unstacked_query_points = tf.unstack(query_points, axis=2, num=2)\n\n        for i, dim in enumerate(index_order):\n            with tf.name_scope(\"dim-\" + str(dim)):\n                queries = unstacked_query_points[dim]\n\n                size_in_indexing_dimension = grid_shape[i + 1]\n\n                # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1\n                # is still a valid index into the grid.\n                max_floor = tf.cast(size_in_indexing_dimension - 2, query_type)\n                min_floor = tf.constant(0.0, dtype=query_type)\n                floor = tf.math.minimum(\n                    tf.math.maximum(min_floor, tf.math.floor(queries)), max_floor\n                )\n                int_floor = tf.cast(floor, tf.dtypes.int32)\n                floors.append(int_floor)\n                ceil = int_floor + 1\n                ceils.append(ceil)\n\n                # alpha has the same type as the grid, as we will directly use alpha\n                # when taking linear combinations of pixel values from the image.\n                alpha = tf.cast(queries - floor, grid_type)\n                min_alpha = tf.constant(0.0, dtype=grid_type)\n                max_alpha = tf.constant(1.0, dtype=grid_type)\n                alpha = tf.math.minimum(tf.math.maximum(min_alpha, alpha), max_alpha)\n\n                # Expand alpha to [b, n, 1] so we can use broadcasting\n                # (since the alpha values don't depend on the channel).\n                alpha = tf.expand_dims(alpha, 2)\n                alphas.append(alpha)\n\n            flattened_grid = tf.reshape(grid, [batch_size * height * width, channels])\n            batch_offsets = tf.reshape(\n                tf.range(batch_size) * height * width, [batch_size, 1]\n            )\n\n        # This wraps tf.gather. We reshape the image data such that the\n        # batch, y, and x coordinates are pulled into the first dimension.\n        # Then we gather. Finally, we reshape the output back. It's possible this\n        # code would be made simpler by using tf.gather_nd.\n        def gather(y_coords, x_coords, name):\n            with tf.name_scope(\"gather-\" + name):\n                linear_coordinates = batch_offsets + y_coords * width + x_coords\n                gathered_values = tf.gather(flattened_grid, linear_coordinates)\n                return tf.reshape(gathered_values, [batch_size, num_queries, channels])\n\n        # grab the pixel values in the 4 corners around each query point\n        top_left = gather(floors[0], floors[1], \"top_left\")\n        top_right = gather(floors[0], ceils[1], \"top_right\")\n        bottom_left = gather(ceils[0], floors[1], \"bottom_left\")\n        bottom_right = gather(ceils[0], ceils[1], \"bottom_right\")\n\n        # now, do the actual interpolation\n        with tf.name_scope(\"interpolate\"):\n            interp_top = alphas[1] * (top_right - top_left) + top_left\n            interp_bottom = alphas[1] * (bottom_right - bottom_left) + bottom_left\n            interp = alphas[0] * (interp_bottom - interp_top) + interp_top\n\n        return interp\n\n\n@tf.function\ndef dense_image_warp(\n    image: types.TensorLike, flow: types.TensorLike, name: Optional[str] = None\n) -> tf.Tensor:\n    \"\"\"Image warping using per-pixel flow vectors.\n\n    Apply a non-linear warp to the image, where the warp is specified by a\n    dense flow field of offset vectors that define the correspondences of\n    pixel values in the output image back to locations in the source image.\n    Specifically, the pixel value at `output[b, j, i, c]` is\n    `images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c]`.\n\n    The locations specified by this formula do not necessarily map to an int\n    index. Therefore, the pixel value is obtained by bilinear\n    interpolation of the 4 nearest pixels around\n    `(b, j - flow[b, j, i, 0], i - flow[b, j, i, 1])`. For locations outside\n    of the image, we use the nearest pixel values at the image boundary.\n\n    PLEASE NOTE: The definition of the flow field above is different from that\n    of optical flow. This function expects the negative forward flow from\n    output image to source image. Given two images `I_1` and `I_2` and the\n    optical flow `F_12` from `I_1` to `I_2`, the image `I_1` can be\n    reconstructed by `I_1_rec = dense_image_warp(I_2, -F_12)`.\n\n    Args:\n      image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.\n      flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.\n      name: A name for the operation (optional).\n\n      Note that image and flow can be of type tf.half, tf.float32, or\n      tf.float64, and do not necessarily have to be the same type.\n\n    Returns:\n      A 4-D float `Tensor` with shape`[batch, height, width, channels]`\n        and same type as input image.\n\n    Raises:\n      ValueError: if height < 2 or width < 2 or the inputs have the wrong\n        number of dimensions.\n    \"\"\"\n    with tf.name_scope(name or \"dense_image_warp\"):\n        image = tf.convert_to_tensor(image)\n        flow = tf.convert_to_tensor(flow)\n        batch_size, height, width, channels = (\n            tf.shape(image)[0],\n            tf.shape(image)[1],\n            tf.shape(image)[2],\n            tf.shape(image)[3],\n        )\n\n        # The flow is defined on the image grid. Turn the flow into a list of query\n        # points in the grid space.\n        grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))\n        stacked_grid = tf.cast(tf.stack([grid_y, grid_x], axis=2), flow.dtype)\n        batched_grid = tf.expand_dims(stacked_grid, axis=0)\n        query_points_on_grid = batched_grid - flow\n        query_points_flattened = tf.reshape(\n            query_points_on_grid, [batch_size, height * width, 2]\n        )\n        # Compute values at the query points, then reshape the result back to the\n        # image grid.\n        interpolated = interpolate_bilinear(image, query_points_flattened)\n        interpolated = tf.reshape(interpolated, [batch_size, height, width, channels])\n        return interpolated\n", "path": "tensorflow_addons/image/dense_image_warp.py"}], "after_files": [{"content": "from tensorflow_addons.utils.test_utils import (  # noqa: F401\n    maybe_run_functions_eagerly,\n    only_run_functions_eagerly,\n    run_custom_and_py_ops,\n    run_with_mixed_precision_policy,\n    pytest_make_parametrize_id,\n    data_format,\n    set_seeds,\n    pytest_addoption,\n    set_global_variables,\n    pytest_configure,\n    device,\n    pytest_generate_tests,\n    pytest_collection_modifyitems,\n)\n\n# fixtures present in this file will be available\n# when running tests and can be referenced with strings\n# https://docs.pytest.org/en/latest/fixture.html#conftest-py-sharing-fixture-functions\n", "path": "tensorflow_addons/conftest.py"}, {"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Image warping using per-pixel flow vectors.\"\"\"\n\nimport numpy as np\nimport tensorflow as tf\n\nfrom tensorflow_addons.utils import types\nfrom typing import Optional\n\n\n@tf.function\ndef interpolate_bilinear(\n    grid: types.TensorLike,\n    query_points: types.TensorLike,\n    indexing: str = \"ij\",\n    name: Optional[str] = None,\n) -> tf.Tensor:\n    \"\"\"Similar to Matlab's interp2 function.\n\n    Finds values for query points on a grid using bilinear interpolation.\n\n    Args:\n      grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.\n      query_points: a 3-D float `Tensor` of N points with shape\n        `[batch, N, 2]`.\n      indexing: whether the query points are specified as row and column (ij),\n        or Cartesian coordinates (xy).\n      name: a name for the operation (optional).\n\n    Returns:\n      values: a 3-D `Tensor` with shape `[batch, N, channels]`\n\n    Raises:\n      ValueError: if the indexing mode is invalid, or if the shape of the\n        inputs invalid.\n    \"\"\"\n    if indexing != \"ij\" and indexing != \"xy\":\n        raise ValueError(\"Indexing mode must be 'ij' or 'xy'\")\n\n    with tf.name_scope(name or \"interpolate_bilinear\"):\n        grid = tf.convert_to_tensor(grid)\n        query_points = tf.convert_to_tensor(query_points)\n\n        # grid shape checks\n        grid_static_shape = grid.shape\n        grid_shape = tf.shape(grid)\n        if grid_static_shape.dims is not None:\n            if len(grid_static_shape) != 4:\n                raise ValueError(\"Grid must be 4D Tensor\")\n            if grid_static_shape[1] is not None and grid_static_shape[1] < 2:\n                raise ValueError(\"Grid height must be at least 2.\")\n            if grid_static_shape[2] is not None and grid_static_shape[2] < 2:\n                raise ValueError(\"Grid width must be at least 2.\")\n        else:\n            with tf.control_dependencies(\n                [\n                    tf.debugging.assert_greater_equal(\n                        grid_shape[1], 2, message=\"Grid height must be at least 2.\"\n                    ),\n                    tf.debugging.assert_greater_equal(\n                        grid_shape[2], 2, message=\"Grid width must be at least 2.\"\n                    ),\n                    tf.debugging.assert_less_equal(\n                        tf.cast(\n                            grid_shape[0] * grid_shape[1] * grid_shape[2],\n                            dtype=tf.dtypes.float32,\n                        ),\n                        np.iinfo(np.int32).max / 8.0,\n                        message=\"The image size or batch size is sufficiently \"\n                        \"large that the linearized addresses used by \"\n                        \"tf.gather may exceed the int32 limit.\",\n                    ),\n                ]\n            ):\n                pass\n\n        # query_points shape checks\n        query_static_shape = query_points.shape\n        query_shape = tf.shape(query_points)\n        if query_static_shape.dims is not None:\n            if len(query_static_shape) != 3:\n                raise ValueError(\"Query points must be 3 dimensional.\")\n            query_hw = query_static_shape[2]\n            if query_hw is not None and query_hw != 2:\n                raise ValueError(\"Query points last dimension must be 2.\")\n        else:\n            with tf.control_dependencies(\n                [\n                    tf.debugging.assert_equal(\n                        query_shape[2],\n                        2,\n                        message=\"Query points last dimension must be 2.\",\n                    )\n                ]\n            ):\n                pass\n\n        batch_size, height, width, channels = (\n            grid_shape[0],\n            grid_shape[1],\n            grid_shape[2],\n            grid_shape[3],\n        )\n\n        num_queries = query_shape[1]\n\n        query_type = query_points.dtype\n        grid_type = grid.dtype\n\n        alphas = []\n        floors = []\n        ceils = []\n        index_order = [0, 1] if indexing == \"ij\" else [1, 0]\n        unstacked_query_points = tf.unstack(query_points, axis=2, num=2)\n\n        for i, dim in enumerate(index_order):\n            with tf.name_scope(\"dim-\" + str(dim)):\n                queries = unstacked_query_points[dim]\n\n                size_in_indexing_dimension = grid_shape[i + 1]\n\n                # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1\n                # is still a valid index into the grid.\n                max_floor = tf.cast(size_in_indexing_dimension - 2, query_type)\n                min_floor = tf.constant(0.0, dtype=query_type)\n                floor = tf.math.minimum(\n                    tf.math.maximum(min_floor, tf.math.floor(queries)), max_floor\n                )\n                int_floor = tf.cast(floor, tf.dtypes.int32)\n                floors.append(int_floor)\n                ceil = int_floor + 1\n                ceils.append(ceil)\n\n                # alpha has the same type as the grid, as we will directly use alpha\n                # when taking linear combinations of pixel values from the image.\n                alpha = tf.cast(queries - floor, grid_type)\n                min_alpha = tf.constant(0.0, dtype=grid_type)\n                max_alpha = tf.constant(1.0, dtype=grid_type)\n                alpha = tf.math.minimum(tf.math.maximum(min_alpha, alpha), max_alpha)\n\n                # Expand alpha to [b, n, 1] so we can use broadcasting\n                # (since the alpha values don't depend on the channel).\n                alpha = tf.expand_dims(alpha, 2)\n                alphas.append(alpha)\n\n            flattened_grid = tf.reshape(grid, [batch_size * height * width, channels])\n            batch_offsets = tf.reshape(\n                tf.range(batch_size) * height * width, [batch_size, 1]\n            )\n\n        # This wraps tf.gather. We reshape the image data such that the\n        # batch, y, and x coordinates are pulled into the first dimension.\n        # Then we gather. Finally, we reshape the output back. It's possible this\n        # code would be made simpler by using tf.gather_nd.\n        def gather(y_coords, x_coords, name):\n            with tf.name_scope(\"gather-\" + name):\n                linear_coordinates = batch_offsets + y_coords * width + x_coords\n                gathered_values = tf.gather(flattened_grid, linear_coordinates)\n                return tf.reshape(gathered_values, [batch_size, num_queries, channels])\n\n        # grab the pixel values in the 4 corners around each query point\n        top_left = gather(floors[0], floors[1], \"top_left\")\n        top_right = gather(floors[0], ceils[1], \"top_right\")\n        bottom_left = gather(ceils[0], floors[1], \"bottom_left\")\n        bottom_right = gather(ceils[0], ceils[1], \"bottom_right\")\n\n        # now, do the actual interpolation\n        with tf.name_scope(\"interpolate\"):\n            interp_top = alphas[1] * (top_right - top_left) + top_left\n            interp_bottom = alphas[1] * (bottom_right - bottom_left) + bottom_left\n            interp = alphas[0] * (interp_bottom - interp_top) + interp_top\n\n        return interp\n\n\ndef _get_dim(x, idx):\n    if x.shape.ndims is None:\n        return tf.shape(x)[idx]\n    return x.shape[idx] or tf.shape(x)[idx]\n\n\n@tf.function\ndef dense_image_warp(\n    image: types.TensorLike, flow: types.TensorLike, name: Optional[str] = None\n) -> tf.Tensor:\n    \"\"\"Image warping using per-pixel flow vectors.\n\n    Apply a non-linear warp to the image, where the warp is specified by a\n    dense flow field of offset vectors that define the correspondences of\n    pixel values in the output image back to locations in the source image.\n    Specifically, the pixel value at `output[b, j, i, c]` is\n    `images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c]`.\n\n    The locations specified by this formula do not necessarily map to an int\n    index. Therefore, the pixel value is obtained by bilinear\n    interpolation of the 4 nearest pixels around\n    `(b, j - flow[b, j, i, 0], i - flow[b, j, i, 1])`. For locations outside\n    of the image, we use the nearest pixel values at the image boundary.\n\n    PLEASE NOTE: The definition of the flow field above is different from that\n    of optical flow. This function expects the negative forward flow from\n    output image to source image. Given two images `I_1` and `I_2` and the\n    optical flow `F_12` from `I_1` to `I_2`, the image `I_1` can be\n    reconstructed by `I_1_rec = dense_image_warp(I_2, -F_12)`.\n\n    Args:\n      image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.\n      flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.\n      name: A name for the operation (optional).\n\n      Note that image and flow can be of type tf.half, tf.float32, or\n      tf.float64, and do not necessarily have to be the same type.\n\n    Returns:\n      A 4-D float `Tensor` with shape`[batch, height, width, channels]`\n        and same type as input image.\n\n    Raises:\n      ValueError: if height < 2 or width < 2 or the inputs have the wrong\n        number of dimensions.\n    \"\"\"\n    with tf.name_scope(name or \"dense_image_warp\"):\n        image = tf.convert_to_tensor(image)\n        flow = tf.convert_to_tensor(flow)\n        batch_size, height, width, channels = (\n            _get_dim(image, 0),\n            _get_dim(image, 1),\n            _get_dim(image, 2),\n            _get_dim(image, 3),\n        )\n\n        # The flow is defined on the image grid. Turn the flow into a list of query\n        # points in the grid space.\n        grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))\n        stacked_grid = tf.cast(tf.stack([grid_y, grid_x], axis=2), flow.dtype)\n        batched_grid = tf.expand_dims(stacked_grid, axis=0)\n        query_points_on_grid = batched_grid - flow\n        query_points_flattened = tf.reshape(\n            query_points_on_grid, [batch_size, height * width, 2]\n        )\n        # Compute values at the query points, then reshape the result back to the\n        # image grid.\n        interpolated = interpolate_bilinear(image, query_points_flattened)\n        interpolated = tf.reshape(interpolated, [batch_size, height, width, channels])\n        return interpolated\n", "path": "tensorflow_addons/image/dense_image_warp.py"}]}