Repository: tensorflow/addons

Issue:
Channel Found None in MaxUnPooling2D
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Colab
- TensorFlow version and how it was installed (source or binary): 2.4.0
- TensorFlow-Addons version and how it was installed (source or binary):
- Python version: 3.7
- Is GPU used? (yes/no):  yes

**Describe the bug**

The channel dimension is None after MaxUnpooling2D layer

**Code to reproduce the issue**

```
def test():
    inputs = tf.keras.layers.Input((None, None, 1))
    x,indices = tf.nn.max_pool_with_argmax(inputs, 2, 2, padding="SAME")
    x = tfa.layers.MaxUnpooling2D()(x,indices)
    x = tf.keras.layers.Conv2D(64, 3, padding="same")(x)
    return tf.keras.Model(inputs, x)
model = test()
model(tf.random.normal([1,128,128,1]))
```

**Other info / logs**

<ipython-input-25-675633950ad9> in test()
      3     x,indices = tf.nn.max_pool_with_argmax(inputs, 2, 2, padding="SAME")
      4     x = tfa.layers.MaxUnpooling2D()(x,indices)
----> 5     x = tf.keras.layers.Conv2D(64, 3, padding="same")(x)
      6     return tf.keras.Model(inputs, x)
      7 model = test()

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
    950     if _in_functional_construction_mode(self, inputs, args, kwargs, input_list):
    951       return self._functional_construction_call(inputs, args, kwargs,
--> 952                                                 input_list)
    953 
    954     # Maintains info about the `Layer.call` stack.

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _functional_construction_call(self, inputs, args, kwargs, input_list)
   1089         # Check input assumptions set after layer building, e.g. input shape.
   1090         outputs = self._keras_tensor_symbolic_call(
-> 1091             inputs, input_masks, args, kwargs)
   1092 
   1093         if outputs is None:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _keras_tensor_symbolic_call(self, inputs, input_masks, args, kwargs)
    820       return nest.map_structure(keras_tensor.KerasTensor, output_signature)
    821     else:
--> 822       return self._infer_output_signature(inputs, args, kwargs, input_masks)
    823 
    824   def _infer_output_signature(self, inputs, args, kwargs, input_masks):

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _infer_output_signature(self, inputs, args, kwargs, input_masks)
    860           # overridden).
    861           # TODO(kaftan): do we maybe_build here, or have we already done it?
--> 862           self._maybe_build(inputs)
    863           outputs = call_fn(inputs, *args, **kwargs)
    864 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _maybe_build(self, inputs)
   2708         # operations.
   2709         with tf_utils.maybe_init_scope(self):
-> 2710           self.build(input_shapes)  # pylint:disable=not-callable
   2711       # We must set also ensure that the layer is marked as built, and the build
   2712       # shape is stored since user defined build functions may not be calling

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/convolutional.py in build(self, input_shape)
    186   def build(self, input_shape):
    187     input_shape = tensor_shape.TensorShape(input_shape)
--> 188     input_channel = self._get_input_channel(input_shape)
    189     if input_channel % self.groups != 0:
    190       raise ValueError(

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/convolutional.py in _get_input_channel(self, input_shape)
    358     channel_axis = self._get_channel_axis()
    359     if input_shape.dims[channel_axis].value is None:
--> 360       raise ValueError('The channel dimension of the inputs '
    361                        'should be defined. Found `None`.')
    362     return int(input_shape[channel_axis])

ValueError: The channel dimension of the inputs should be defined. Found `None`.

Discussion:
WindQAQ: /cc @thaink for visibility.
bot-of-gabrieldemarmiesse: @thaink 

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.


Patch:
diff --git a/tensorflow_addons/layers/max_unpooling_2d.py b/tensorflow_addons/layers/max_unpooling_2d.py
index c8e1bf3..4e83fd1 100644
--- a/tensorflow_addons/layers/max_unpooling_2d.py
+++ b/tensorflow_addons/layers/max_unpooling_2d.py
@@ -59,6 +59,7 @@ def _max_unpooling_2d(updates, mask, pool_size=(2, 2), strides=(2, 2), padding="
     def func(updates, mask):
         mask = tf.cast(mask, "int32")
         input_shape = tf.shape(updates, out_type="int32")
+        input_shape = [updates.shape[i] or input_shape[i] for i in range(4)]
         output_shape = _calculate_output_shape(input_shape, pool_size, strides, padding)
 
         # Calculates indices for batch, height, width and feature maps.


Test patch:
diff --git a/tensorflow_addons/layers/tests/max_unpooling_2d_test.py b/tensorflow_addons/layers/tests/max_unpooling_2d_test.py
index ce4a300..7838c67 100644
--- a/tensorflow_addons/layers/tests/max_unpooling_2d_test.py
+++ b/tensorflow_addons/layers/tests/max_unpooling_2d_test.py
@@ -238,3 +238,15 @@ def test_padding_valid_complex():
 
     output = MaxUnpooling2D(pool_size=(2, 3), padding="VALID")(updates, indices).numpy()
     np.testing.assert_array_equal(expected_output, output)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_symbolic_tensor_shape():
+    valid_input = tf.keras.layers.Input((None, None, 1))
+    updates, indices = tf.nn.max_pool_with_argmax(
+        valid_input, ksize=[2, 2], strides=[2, 2], padding="SAME"
+    )
+    output = MaxUnpooling2D(pool_size=(2, 2), strides=(2, 2), padding="SAME")(
+        updates, indices
+    )
+    np.testing.assert_array_equal(valid_input.shape.as_list(), output.shape.as_list())