Repository: tensorflow/addons

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.



Discussion:
It's because you are using functional API of keras. In this case, `image` is a symbolic tensor, and thus, the reshape in [this line](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/dense_image_warp.py#L251) based on `tf.shape` won't provide more info.

One possible fix is that you can use `.set_shape(image.shape)` after the `tfa.image_dense_image_warp(image, flow)`.

```python
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)
warped.set_shape(image.shape)
```
Thank you very much for the suggestion. The possible fix solves the issue, although I feel like it could be beneficial to modify the function to adapt to this usage. It is normal to use placeholder input with `None` batch size when compiling a model, and it would be cleaner if we don't have to manually set the output shape.
> Thank you very much for the suggestion. The possible fix solves the issue, although I feel like it could be beneficial to modify the function to adapt to this usage. It is normal to use placeholder input with `None` batch size when compiling a model, and it would be cleaner if we don't have to manually set the output shape.

Sure, we are more than welcome to have a PR for this. Thanks for the suggestion!

Resolving patch:
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 @@ def gather(y_coords, x_coords, name):
         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 @@ def dense_image_warp(
         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


Test patch:
diff --git a/tensorflow_addons/image/tests/dense_image_warp_test.py b/tensorflow_addons/image/tests/dense_image_warp_test.py
--- a/tensorflow_addons/image/tests/dense_image_warp_test.py
+++ b/tensorflow_addons/image/tests/dense_image_warp_test.py
@@ -251,3 +251,11 @@ 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])
diff --git a/tensorflow_addons/utils/test_utils.py b/tensorflow_addons/utils/test_utils.py
--- a/tensorflow_addons/utils/test_utils.py
+++ b/tensorflow_addons/utils/test_utils.py
@@ -88,6 +88,12 @@ def maybe_run_functions_eagerly(request):
     request.addfinalizer(finalizer)
 
 
+@pytest.fixture(scope="function")
+def only_run_functions_eagerly(request):
+    tf.config.experimental_run_functions_eagerly(True)
+    request.addfinalizer(finalizer)
+
+
 @pytest.fixture(scope="function", params=["custom_ops", "py_ops"])
 def run_custom_and_py_ops(request):
     previous_py_ops_value = options.TF_ADDONS_PY_OPS