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:
Thanks for reporting. Some static values are missing while slicing. Can you try to insert the following line

```python
input_shape = [updates.shape[i] or input_shape[i] for i in range(4)]
```

under

https://github.com/tensorflow/addons/blob/445638b02f7cdeb2d33b2cda1b66fb8a4f639366/tensorflow_addons/layers/max_unpooling_2d.py#L61
@WindQAQ  ok, I will try and revert it back here, however the problem is solved when I constructing the model by subclass the tf.keras.Model instead of using keras functional API.

But this raises another problem which the the tf function retracing
I think the @tf.function should be removed inside this layer as when we annotate our training code with @tf.function, this will become nested tf function
@WindQAQ  it works after adding that line of code

`tf.function` here is used for tflite conversion. Hi @thaink, can you check if the solution is valid? Thank you!
@WindQAQ but the tf.function here will cause retracing issue during training
The solution `input_shape = [updates.shape[i] or input_shape[i] for i in range(4)]` is valid.

`tf.function` is widely used. As you can see in https://github.com/tensorflow/addons/tree/master/tensorflow_addons/image, it is used for a lot of functionalities. I wonder what is the benefit of annotating your training with @tf.function? It seems to block you from using fused operations and tf.function is only significantly faster with "graphs with many small ops".


@thaink I did a benchmark before and the speed with tf function is faster compared to without tf function

Resolving patch:
diff --git a/tensorflow_addons/layers/max_unpooling_2d.py b/tensorflow_addons/layers/max_unpooling_2d.py
--- 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
--- 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())