We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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`.

--- 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/layers/max_unpooling_2d.py`
Content:
```
1 # Copyright 2020 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 """MaxUnpooling2D operation."""
16 
17 import tensorflow as tf
18 
19 from typeguard import typechecked
20 from typing import Union
21 
22 from tensorflow_addons.utils.keras_utils import normalize_tuple
23 
24 
25 def _calculate_output_shape(input_shape, pool_size, strides, padding):
26     """Calculates the shape of the unpooled output."""
27     if padding == "VALID":
28         output_shape = (
29             input_shape[0],
30             (input_shape[1] - 1) * strides[0] + pool_size[0],
31             (input_shape[2] - 1) * strides[1] + pool_size[1],
32             input_shape[3],
33         )
34     elif padding == "SAME":
35         output_shape = (
36             input_shape[0],
37             input_shape[1] * strides[0],
38             input_shape[2] * strides[1],
39             input_shape[3],
40         )
41     else:
42         raise ValueError('Padding must be a string from: "SAME", "VALID"')
43     return output_shape
44 
45 
46 def _max_unpooling_2d(updates, mask, pool_size=(2, 2), strides=(2, 2), padding="SAME"):
47     """Unpool the outputs of a maximum pooling operation."""
48     pool_size_attr = " ".join(["i: %d" % v for v in pool_size])
49     strides_attr = " ".join(["i: %d" % v for v in strides])
50     experimental_implements = [
51         'name: "addons:MaxUnpooling2D"',
52         'attr { key: "pool_size" value { list {%s} } }' % pool_size_attr,
53         'attr { key: "strides" value { list {%s} } }' % strides_attr,
54         'attr { key: "padding" value { s: "%s" } }' % padding,
55     ]
56     experimental_implements = " ".join(experimental_implements)
57 
58     @tf.function(experimental_implements=experimental_implements)
59     def func(updates, mask):
60         mask = tf.cast(mask, "int32")
61         input_shape = tf.shape(updates, out_type="int32")
62         output_shape = _calculate_output_shape(input_shape, pool_size, strides, padding)
63 
64         # Calculates indices for batch, height, width and feature maps.
65         one_like_mask = tf.ones_like(mask, dtype="int32")
66         batch_shape = tf.concat([[input_shape[0]], [1], [1], [1]], axis=0)
67         batch_range = tf.reshape(
68             tf.range(output_shape[0], dtype="int32"), shape=batch_shape
69         )
70         b = one_like_mask * batch_range
71         y = mask // (output_shape[2] * output_shape[3])
72         x = (mask // output_shape[3]) % output_shape[2]
73         feature_range = tf.range(output_shape[3], dtype="int32")
74         f = one_like_mask * feature_range
75 
76         # Transposes indices & reshape update values to one dimension.
77         updates_size = tf.size(updates)
78         indices = tf.transpose(tf.reshape(tf.stack([b, y, x, f]), [4, updates_size]))
79         values = tf.reshape(updates, [updates_size])
80         ret = tf.scatter_nd(indices, values, output_shape)
81         return ret
82 
83     return func(updates, mask)
84 
85 
86 @tf.keras.utils.register_keras_serializable(package="Addons")
87 class MaxUnpooling2D(tf.keras.layers.Layer):
88     """Unpool the outputs of a maximum pooling operation.
89 
90     This function currently does not support outputs of MaxPoolingWithArgMax in
91     following cases:
92     - include_batch_in_index equals true.
93     - input_shape is not divisible by strides if padding is "SAME".
94     - (input_shape - pool_size) is not divisible by strides if padding is "VALID".
95     - The max pooling operation results in duplicate values in updates and mask.
96 
97     Args:
98       updates: The pooling result from max pooling.
99       mask: the argmax result corresponds to above max values.
100       pool_size: The filter that max pooling was performed with. Default: (2, 2).
101       strides: The strides that max pooling was performed with. Default: (2, 2).
102       padding: The padding that max pooling was performed with. Default: "SAME".
103     Input shape:
104       4D tensor with shape: `(batch_size, height, width, channel)`.
105     Output shape:
106       4D tensor with the same shape as the input of max pooling operation.
107     """
108 
109     @typechecked
110     def __init__(
111         self,
112         pool_size: Union[list, tuple, int] = (2, 2),
113         strides: Union[list, tuple, int] = (2, 2),
114         padding: str = "SAME",
115         **kwargs,
116     ):
117         super(MaxUnpooling2D, self).__init__(**kwargs)
118 
119         if padding != "SAME" and padding != "VALID":
120             raise ValueError('Padding must be a string from: "SAME", "VALID"')
121 
122         self.pool_size = normalize_tuple(pool_size, 2, "pool_size")
123         self.strides = normalize_tuple(strides, 2, "strides")
124         self.padding = padding
125 
126     def call(self, updates, mask):
127         return _max_unpooling_2d(
128             updates,
129             mask,
130             pool_size=self.pool_size,
131             strides=self.strides,
132             padding=self.padding,
133         )
134 
135     def compute_output_shape(self, input_shapes):
136         input_shape = input_shapes[1]
137         return _calculate_output_shape(
138             input_shape, self.pool_size, self.strides, self.padding
139         )
140 
141     def get_config(self):
142         config = super(MaxUnpooling2D, self).get_config()
143         config["pool_size"] = self.pool_size
144         config["strides"] = self.strides
145         config["padding"] = self.padding
146         return config
147
```
--- 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/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 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.

{"golden_diff": "diff --git a/tensorflow_addons/layers/max_unpooling_2d.py b/tensorflow_addons/layers/max_unpooling_2d.py\n--- a/tensorflow_addons/layers/max_unpooling_2d.py\n+++ b/tensorflow_addons/layers/max_unpooling_2d.py\n@@ -59,6 +59,7 @@\n     def func(updates, mask):\n         mask = tf.cast(mask, \"int32\")\n         input_shape = tf.shape(updates, out_type=\"int32\")\n+        input_shape = [updates.shape[i] or input_shape[i] for i in range(4)]\n         output_shape = _calculate_output_shape(input_shape, pool_size, strides, padding)\n \n         # Calculates indices for batch, height, width and feature maps.\n", "issue": "Channel Found None in MaxUnPooling2D\n**System information**\r\n- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Colab\r\n- TensorFlow version and how it was installed (source or binary): 2.4.0\r\n- TensorFlow-Addons version and how it was installed (source or binary):\r\n- Python version: 3.7\r\n- Is GPU used? (yes/no):  yes\r\n\r\n**Describe the bug**\r\n\r\nThe channel dimension is None after MaxUnpooling2D layer\r\n\r\n**Code to reproduce the issue**\r\n\r\n```\r\ndef test():\r\n    inputs = tf.keras.layers.Input((None, None, 1))\r\n    x,indices = tf.nn.max_pool_with_argmax(inputs, 2, 2, padding=\"SAME\")\r\n    x = tfa.layers.MaxUnpooling2D()(x,indices)\r\n    x = tf.keras.layers.Conv2D(64, 3, padding=\"same\")(x)\r\n    return tf.keras.Model(inputs, x)\r\nmodel = test()\r\nmodel(tf.random.normal([1,128,128,1]))\r\n```\r\n\r\n**Other info / logs**\r\n\r\n<ipython-input-25-675633950ad9> in test()\r\n      3     x,indices = tf.nn.max_pool_with_argmax(inputs, 2, 2, padding=\"SAME\")\r\n      4     x = tfa.layers.MaxUnpooling2D()(x,indices)\r\n----> 5     x = tf.keras.layers.Conv2D(64, 3, padding=\"same\")(x)\r\n      6     return tf.keras.Model(inputs, x)\r\n      7 model = test()\r\n\r\n/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)\r\n    950     if _in_functional_construction_mode(self, inputs, args, kwargs, input_list):\r\n    951       return self._functional_construction_call(inputs, args, kwargs,\r\n--> 952                                                 input_list)\r\n    953 \r\n    954     # Maintains info about the `Layer.call` stack.\r\n\r\n/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _functional_construction_call(self, inputs, args, kwargs, input_list)\r\n   1089         # Check input assumptions set after layer building, e.g. input shape.\r\n   1090         outputs = self._keras_tensor_symbolic_call(\r\n-> 1091             inputs, input_masks, args, kwargs)\r\n   1092 \r\n   1093         if outputs is None:\r\n\r\n/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)\r\n    820       return nest.map_structure(keras_tensor.KerasTensor, output_signature)\r\n    821     else:\r\n--> 822       return self._infer_output_signature(inputs, args, kwargs, input_masks)\r\n    823 \r\n    824   def _infer_output_signature(self, inputs, args, kwargs, input_masks):\r\n\r\n/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _infer_output_signature(self, inputs, args, kwargs, input_masks)\r\n    860           # overridden).\r\n    861           # TODO(kaftan): do we maybe_build here, or have we already done it?\r\n--> 862           self._maybe_build(inputs)\r\n    863           outputs = call_fn(inputs, *args, **kwargs)\r\n    864 \r\n\r\n/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _maybe_build(self, inputs)\r\n   2708         # operations.\r\n   2709         with tf_utils.maybe_init_scope(self):\r\n-> 2710           self.build(input_shapes)  # pylint:disable=not-callable\r\n   2711       # We must set also ensure that the layer is marked as built, and the build\r\n   2712       # shape is stored since user defined build functions may not be calling\r\n\r\n/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/convolutional.py in build(self, input_shape)\r\n    186   def build(self, input_shape):\r\n    187     input_shape = tensor_shape.TensorShape(input_shape)\r\n--> 188     input_channel = self._get_input_channel(input_shape)\r\n    189     if input_channel % self.groups != 0:\r\n    190       raise ValueError(\r\n\r\n/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/convolutional.py in _get_input_channel(self, input_shape)\r\n    358     channel_axis = self._get_channel_axis()\r\n    359     if input_shape.dims[channel_axis].value is None:\r\n--> 360       raise ValueError('The channel dimension of the inputs '\r\n    361                        'should be defined. Found `None`.')\r\n    362     return int(input_shape[channel_axis])\r\n\r\nValueError: The channel dimension of the inputs should be defined. Found `None`.\n", "before_files": [{"content": "# Copyright 2020 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\"\"\"MaxUnpooling2D operation.\"\"\"\n\nimport tensorflow as tf\n\nfrom typeguard import typechecked\nfrom typing import Union\n\nfrom tensorflow_addons.utils.keras_utils import normalize_tuple\n\n\ndef _calculate_output_shape(input_shape, pool_size, strides, padding):\n    \"\"\"Calculates the shape of the unpooled output.\"\"\"\n    if padding == \"VALID\":\n        output_shape = (\n            input_shape[0],\n            (input_shape[1] - 1) * strides[0] + pool_size[0],\n            (input_shape[2] - 1) * strides[1] + pool_size[1],\n            input_shape[3],\n        )\n    elif padding == \"SAME\":\n        output_shape = (\n            input_shape[0],\n            input_shape[1] * strides[0],\n            input_shape[2] * strides[1],\n            input_shape[3],\n        )\n    else:\n        raise ValueError('Padding must be a string from: \"SAME\", \"VALID\"')\n    return output_shape\n\n\ndef _max_unpooling_2d(updates, mask, pool_size=(2, 2), strides=(2, 2), padding=\"SAME\"):\n    \"\"\"Unpool the outputs of a maximum pooling operation.\"\"\"\n    pool_size_attr = \" \".join([\"i: %d\" % v for v in pool_size])\n    strides_attr = \" \".join([\"i: %d\" % v for v in strides])\n    experimental_implements = [\n        'name: \"addons:MaxUnpooling2D\"',\n        'attr { key: \"pool_size\" value { list {%s} } }' % pool_size_attr,\n        'attr { key: \"strides\" value { list {%s} } }' % strides_attr,\n        'attr { key: \"padding\" value { s: \"%s\" } }' % padding,\n    ]\n    experimental_implements = \" \".join(experimental_implements)\n\n    @tf.function(experimental_implements=experimental_implements)\n    def func(updates, mask):\n        mask = tf.cast(mask, \"int32\")\n        input_shape = tf.shape(updates, out_type=\"int32\")\n        output_shape = _calculate_output_shape(input_shape, pool_size, strides, padding)\n\n        # Calculates indices for batch, height, width and feature maps.\n        one_like_mask = tf.ones_like(mask, dtype=\"int32\")\n        batch_shape = tf.concat([[input_shape[0]], [1], [1], [1]], axis=0)\n        batch_range = tf.reshape(\n            tf.range(output_shape[0], dtype=\"int32\"), shape=batch_shape\n        )\n        b = one_like_mask * batch_range\n        y = mask // (output_shape[2] * output_shape[3])\n        x = (mask // output_shape[3]) % output_shape[2]\n        feature_range = tf.range(output_shape[3], dtype=\"int32\")\n        f = one_like_mask * feature_range\n\n        # Transposes indices & reshape update values to one dimension.\n        updates_size = tf.size(updates)\n        indices = tf.transpose(tf.reshape(tf.stack([b, y, x, f]), [4, updates_size]))\n        values = tf.reshape(updates, [updates_size])\n        ret = tf.scatter_nd(indices, values, output_shape)\n        return ret\n\n    return func(updates, mask)\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass MaxUnpooling2D(tf.keras.layers.Layer):\n    \"\"\"Unpool the outputs of a maximum pooling operation.\n\n    This function currently does not support outputs of MaxPoolingWithArgMax in\n    following cases:\n    - include_batch_in_index equals true.\n    - input_shape is not divisible by strides if padding is \"SAME\".\n    - (input_shape - pool_size) is not divisible by strides if padding is \"VALID\".\n    - The max pooling operation results in duplicate values in updates and mask.\n\n    Args:\n      updates: The pooling result from max pooling.\n      mask: the argmax result corresponds to above max values.\n      pool_size: The filter that max pooling was performed with. Default: (2, 2).\n      strides: The strides that max pooling was performed with. Default: (2, 2).\n      padding: The padding that max pooling was performed with. Default: \"SAME\".\n    Input shape:\n      4D tensor with shape: `(batch_size, height, width, channel)`.\n    Output shape:\n      4D tensor with the same shape as the input of max pooling operation.\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        pool_size: Union[list, tuple, int] = (2, 2),\n        strides: Union[list, tuple, int] = (2, 2),\n        padding: str = \"SAME\",\n        **kwargs,\n    ):\n        super(MaxUnpooling2D, self).__init__(**kwargs)\n\n        if padding != \"SAME\" and padding != \"VALID\":\n            raise ValueError('Padding must be a string from: \"SAME\", \"VALID\"')\n\n        self.pool_size = normalize_tuple(pool_size, 2, \"pool_size\")\n        self.strides = normalize_tuple(strides, 2, \"strides\")\n        self.padding = padding\n\n    def call(self, updates, mask):\n        return _max_unpooling_2d(\n            updates,\n            mask,\n            pool_size=self.pool_size,\n            strides=self.strides,\n            padding=self.padding,\n        )\n\n    def compute_output_shape(self, input_shapes):\n        input_shape = input_shapes[1]\n        return _calculate_output_shape(\n            input_shape, self.pool_size, self.strides, self.padding\n        )\n\n    def get_config(self):\n        config = super(MaxUnpooling2D, self).get_config()\n        config[\"pool_size\"] = self.pool_size\n        config[\"strides\"] = self.strides\n        config[\"padding\"] = self.padding\n        return config\n", "path": "tensorflow_addons/layers/max_unpooling_2d.py"}], "after_files": [{"content": "# Copyright 2020 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\"\"\"MaxUnpooling2D operation.\"\"\"\n\nimport tensorflow as tf\n\nfrom typeguard import typechecked\nfrom typing import Union\n\nfrom tensorflow_addons.utils.keras_utils import normalize_tuple\n\n\ndef _calculate_output_shape(input_shape, pool_size, strides, padding):\n    \"\"\"Calculates the shape of the unpooled output.\"\"\"\n    if padding == \"VALID\":\n        output_shape = (\n            input_shape[0],\n            (input_shape[1] - 1) * strides[0] + pool_size[0],\n            (input_shape[2] - 1) * strides[1] + pool_size[1],\n            input_shape[3],\n        )\n    elif padding == \"SAME\":\n        output_shape = (\n            input_shape[0],\n            input_shape[1] * strides[0],\n            input_shape[2] * strides[1],\n            input_shape[3],\n        )\n    else:\n        raise ValueError('Padding must be a string from: \"SAME\", \"VALID\"')\n    return output_shape\n\n\ndef _max_unpooling_2d(updates, mask, pool_size=(2, 2), strides=(2, 2), padding=\"SAME\"):\n    \"\"\"Unpool the outputs of a maximum pooling operation.\"\"\"\n    pool_size_attr = \" \".join([\"i: %d\" % v for v in pool_size])\n    strides_attr = \" \".join([\"i: %d\" % v for v in strides])\n    experimental_implements = [\n        'name: \"addons:MaxUnpooling2D\"',\n        'attr { key: \"pool_size\" value { list {%s} } }' % pool_size_attr,\n        'attr { key: \"strides\" value { list {%s} } }' % strides_attr,\n        'attr { key: \"padding\" value { s: \"%s\" } }' % padding,\n    ]\n    experimental_implements = \" \".join(experimental_implements)\n\n    @tf.function(experimental_implements=experimental_implements)\n    def func(updates, mask):\n        mask = tf.cast(mask, \"int32\")\n        input_shape = tf.shape(updates, out_type=\"int32\")\n        input_shape = [updates.shape[i] or input_shape[i] for i in range(4)]\n        output_shape = _calculate_output_shape(input_shape, pool_size, strides, padding)\n\n        # Calculates indices for batch, height, width and feature maps.\n        one_like_mask = tf.ones_like(mask, dtype=\"int32\")\n        batch_shape = tf.concat([[input_shape[0]], [1], [1], [1]], axis=0)\n        batch_range = tf.reshape(\n            tf.range(output_shape[0], dtype=\"int32\"), shape=batch_shape\n        )\n        b = one_like_mask * batch_range\n        y = mask // (output_shape[2] * output_shape[3])\n        x = (mask // output_shape[3]) % output_shape[2]\n        feature_range = tf.range(output_shape[3], dtype=\"int32\")\n        f = one_like_mask * feature_range\n\n        # Transposes indices & reshape update values to one dimension.\n        updates_size = tf.size(updates)\n        indices = tf.transpose(tf.reshape(tf.stack([b, y, x, f]), [4, updates_size]))\n        values = tf.reshape(updates, [updates_size])\n        ret = tf.scatter_nd(indices, values, output_shape)\n        return ret\n\n    return func(updates, mask)\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass MaxUnpooling2D(tf.keras.layers.Layer):\n    \"\"\"Unpool the outputs of a maximum pooling operation.\n\n    This function currently does not support outputs of MaxPoolingWithArgMax in\n    following cases:\n    - include_batch_in_index equals true.\n    - input_shape is not divisible by strides if padding is \"SAME\".\n    - (input_shape - pool_size) is not divisible by strides if padding is \"VALID\".\n    - The max pooling operation results in duplicate values in updates and mask.\n\n    Args:\n      updates: The pooling result from max pooling.\n      mask: the argmax result corresponds to above max values.\n      pool_size: The filter that max pooling was performed with. Default: (2, 2).\n      strides: The strides that max pooling was performed with. Default: (2, 2).\n      padding: The padding that max pooling was performed with. Default: \"SAME\".\n    Input shape:\n      4D tensor with shape: `(batch_size, height, width, channel)`.\n    Output shape:\n      4D tensor with the same shape as the input of max pooling operation.\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        pool_size: Union[list, tuple, int] = (2, 2),\n        strides: Union[list, tuple, int] = (2, 2),\n        padding: str = \"SAME\",\n        **kwargs,\n    ):\n        super(MaxUnpooling2D, self).__init__(**kwargs)\n\n        if padding != \"SAME\" and padding != \"VALID\":\n            raise ValueError('Padding must be a string from: \"SAME\", \"VALID\"')\n\n        self.pool_size = normalize_tuple(pool_size, 2, \"pool_size\")\n        self.strides = normalize_tuple(strides, 2, \"strides\")\n        self.padding = padding\n\n    def call(self, updates, mask):\n        return _max_unpooling_2d(\n            updates,\n            mask,\n            pool_size=self.pool_size,\n            strides=self.strides,\n            padding=self.padding,\n        )\n\n    def compute_output_shape(self, input_shapes):\n        input_shape = input_shapes[1]\n        return _calculate_output_shape(\n            input_shape, self.pool_size, self.strides, self.padding\n        )\n\n    def get_config(self):\n        config = super(MaxUnpooling2D, self).get_config()\n        config[\"pool_size\"] = self.pool_size\n        config[\"strides\"] = self.strides\n        config[\"padding\"] = self.padding\n        return config\n", "path": "tensorflow_addons/layers/max_unpooling_2d.py"}]}