Time distributed WeightNormalization result in incompatible shape

[asamax]

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

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]

[seanpmorgan]

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