Repository: tensorflow/addons

Issue:
Time distributed WeightNormalization result in incompatible shape
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04.4 LTS
- TensorFlow version and how it was installed (source or binary): tensorflow-gpu 2.0.0, binary
- TensorFlow-Addons version and how it was installed (source or binary): tensorflow-addons 0.6.0, binary
- Python version: Python 3.6.9
- Is GPU used? (yes/no): yes

**Describe the bug**
Time distributed WeightNormalization result in incompatible shape error

Error occur in `input_spec.assert_input_compatibility(self.input_spec, inputs, self.name)`
 - WeightNormalization ON: `self.input_spec (shape=[32, 64, 64, 3], ndim=4)`
 - WeightNormalization OFF: `self.input_spec (ndim=4, axes={-1:3})`

**Code to reproduce the issue**

```python
import tensorflow as tf
from tensorflow_addons.layers.wrappers import WeightNormalization

weight_norm = True

batch_shape = (32, 16, 64, 64, 3)
inputs = tf.keras.layers.Input(batch_shape=batch_shape)
if weight_norm is True:
    a = tf.keras.layers.Conv2D(3, 5)
    b = WeightNormalization(a)
    outputs = tf.keras.layers.TimeDistributed(b)(inputs)
else:
    outputs = tf.keras.layers.TimeDistributed(tf.keras.layers.Conv2D(3, 5))(inputs)
model = tf.keras.Model(inputs, outputs)
```

**Other info / logs**

```
WARNING:tensorflow:From /opt/anaconda/lib/python3.6/site-packages/tensorflow_addons/layers/wrappers.py:84: Layer.add_variable (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.
Instructions for updating:
Please use `layer.add_weight` method instead.
Traceback (most recent call last):
  File "/workarea/projects/magic/MagIC/applications/temporal/app_test.py", line 11, in <module>
    outputs = tf.keras.layers.TimeDistributed(b)(inputs)
  File "/opt/anaconda/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 842, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "/opt/anaconda/lib/python3.6/site-packages/tensorflow_core/python/keras/layers/wrappers.py", line 254, in call
    y = self.layer(inputs, **kwargs)
  File "/opt/anaconda/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 812, in __call__
    self.name)
  File "/opt/anaconda/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/input_spec.py", line 224, in assert_input_compatibility
    ', found shape=' + str(shape))
ValueError: Input 0 is incompatible with layer weight_normalization: expected shape=[32, 64, 64, 3], found shape=[512, 64, 64, 3]
```


Discussion:
Thanks for reporting! Will need to look into whats happening when two wrappers are used on the layer.

Resolving patch:
diff --git a/tensorflow_addons/layers/wrappers.py b/tensorflow_addons/layers/wrappers.py
--- a/tensorflow_addons/layers/wrappers.py
+++ b/tensorflow_addons/layers/wrappers.py
@@ -62,7 +62,8 @@ def __init__(self, layer, data_init=True, **kwargs):
     def build(self, input_shape):
         """Build `Layer`"""
         input_shape = tf.TensorShape(input_shape).as_list()
-        self.input_spec = tf.keras.layers.InputSpec(shape=input_shape)
+        self.input_spec = tf.keras.layers.InputSpec(
+            shape=[None] + input_shape[1:])
 
         if not self.layer.built:
             self.layer.build(input_shape)


Test patch:
diff --git a/tensorflow_addons/layers/wrappers_test.py b/tensorflow_addons/layers/wrappers_test.py
--- a/tensorflow_addons/layers/wrappers_test.py
+++ b/tensorflow_addons/layers/wrappers_test.py
@@ -73,6 +73,14 @@ def test_weightnorm_non_kernel_layer(self):
             wn_wrapper = wrappers.WeightNormalization(non_kernel_layer)
             wn_wrapper(images)
 
+    def test_weightnorm_with_time_dist(self):
+        batch_shape = (32, 16, 64, 64, 3)
+        inputs = tf.keras.layers.Input(batch_shape=batch_shape)
+        a = tf.keras.layers.Conv2D(3, 5)
+        b = wrappers.WeightNormalization(a)
+        out = tf.keras.layers.TimeDistributed(b)(inputs)
+        model = tf.keras.Model(inputs, out)
+
 
 if __name__ == "__main__":
     tf.test.main()