Repository: tensorflow/addons

Issue:
WeightNormalization with RNNs: shape issue
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab
- TensorFlow version and how it was installed (source or binary): 2.0.0 binary
- TensorFlow-Addons version and how it was installed (source or binary): 0.6.0 binary
- Python version: 3.6.8
- Is GPU used? (yes/no): yes

**Describe the bug**

WeightNormalization layer wrapper cannot be used with RNNs if the input sequence has undetermined length. See code for errors.

**Code to reproduce the issue**

```
import tensorflow as tf
import tensorflow_addons as tfa

n_features = 3
seq_length = None
rnn_units = 4

input_layer = tf.keras.layers.Input(shape=(seq_length, n_features))
rnn_layer = tf.keras.layers.SimpleRNN(rnn_units)
dense_layer = tf.keras.layers.Dense(1)
wn_rnn_layer = tfa.layers.WeightNormalization(rnn_layer)
wn_model = tf.keras.models.Sequential(layers=(input_layer, wn_rnn_layer, dense_layer))
```
yields
```
ValueError: as_list() is not defined on an unknown TensorShape.
```

Note that:
1. The same code without using `WeightNormalization` runs.
2. Interestingly, adding the lines
```
batch_size = 1
input_layer = tf.keras.layers.Input(batch_shape=(batch_size, seq_length, n_features))
rnn_layer = tf.keras.layers.SimpleRNN(rnn_units, return_sequences=True)
dense_layer = tf.keras.layers.Dense(1)
wn_rnn_layer = tfa.layers.WeightNormalization(rnn_layer)
wn_model = tf.keras.models.Sequential(layers=(input_layer, wn_rnn_layer, dense_layer))
```
gives
```
IndexError: list assignment index out of range
```
instead.

**Other info / logs**



Discussion:

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
@@ -58,31 +58,34 @@ def __init__(self, layer, data_init=True, **kwargs):
         super(WeightNormalization, self).__init__(layer, **kwargs)
         self.data_init = data_init
         self._track_trackable(layer, name='layer')
+        self.is_rnn = isinstance(self.layer, tf.keras.layers.RNN)
 
     def build(self, input_shape):
         """Build `Layer`"""
-        input_shape = tf.TensorShape(input_shape).as_list()
+        input_shape = tf.TensorShape(input_shape)
         self.input_spec = tf.keras.layers.InputSpec(
             shape=[None] + input_shape[1:])
 
         if not self.layer.built:
             self.layer.build(input_shape)
 
-        if not hasattr(self.layer, 'kernel'):
+        kernel_layer = self.layer.cell if self.is_rnn else self.layer
+
+        if not hasattr(kernel_layer, 'kernel'):
             raise ValueError('`WeightNormalization` must wrap a layer that'
                              ' contains a `kernel` for weights')
 
         # The kernel's filter or unit dimension is -1
-        self.layer_depth = int(self.layer.kernel.shape[-1])
-        self.kernel_norm_axes = list(range(self.layer.kernel.shape.rank - 1))
+        self.layer_depth = int(kernel_layer.kernel.shape[-1])
+        self.kernel_norm_axes = list(range(kernel_layer.kernel.shape.rank - 1))
 
         self.g = self.add_weight(
             name='g',
             shape=(self.layer_depth,),
             initializer='ones',
-            dtype=self.layer.kernel.dtype,
+            dtype=kernel_layer.kernel.dtype,
             trainable=True)
-        self.v = self.layer.kernel
+        self.v = kernel_layer.kernel
 
         self._initialized = self.add_weight(
             name='initialized',
@@ -100,7 +103,10 @@ def build(self, input_shape):
                     layer_config)
                 self._naked_clone_layer.build(input_shape)
                 self._naked_clone_layer.set_weights(self.layer.get_weights())
-                self._naked_clone_layer.activation = None
+                if self.is_rnn:
+                    self._naked_clone_layer.cell.activation = None
+                else:
+                    self._naked_clone_layer.activation = None
 
         self.built = True
 


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
@@ -89,6 +89,13 @@ def test_weightnorm_with_time_dist(self):
         out = tf.keras.layers.TimeDistributed(b)(inputs)
         model = tf.keras.Model(inputs, out)
 
+    def test_weightnorm_with_rnn(self):
+        inputs = tf.keras.layers.Input(shape=(None, 3))
+        rnn_layer = tf.keras.layers.SimpleRNN(4)
+        wt_rnn = wrappers.WeightNormalization(rnn_layer)
+        dense = tf.keras.layers.Dense(1)
+        model = tf.keras.models.Sequential(layers=[inputs, wt_rnn, dense])
+
     def test_save_file_h5(self):
         self.create_tempfile('wrapper_test_model.h5')
         conv = tf.keras.layers.Conv1D(1, 1)