Hotfix: Compile Keras WN models without tf_function (#429)

Author: seanpmorgan

tensorflow_addons/layers/wrappers.py

@@ -88,7 +88,6 @@ def build(self, input_shape):
 
         super(WeightNormalization, self).build()
 
-    @tf.function
     def call(self, inputs):
         """Call `Layer`"""
         if not self._initialized:
@@ -131,6 +130,7 @@ def _init_norm(self):
             self.g.assign(
                 tf.reshape(tf.linalg.norm(flat, axis=0), (self.layer_depth,)))
 
+    # TODO: Get data init to work with tf_function compile #428
     def _data_dep_init(self, inputs):
         """Data dependent initialization."""
 

tensorflow_addons/layers/wrappers_test.py

@@ -26,14 +26,16 @@
 
 @test_utils.run_all_in_graph_and_eager_modes
 class WeightNormalizationTest(tf.test.TestCase):
+    # TODO: Get data init to work with tf_function compile #428
     def test_weightnorm_dense_train(self):
         model = tf.keras.models.Sequential()
         model.add(
             wrappers.WeightNormalization(
                 tf.keras.layers.Dense(2), input_shape=(3, 4)))
         model.compile(
             optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.001),
-            loss='mse')
+            loss='mse',
+            experimental_run_tf_function=False)
         model.fit(
             np.random.random((10, 3, 4)),
             np.random.random((10, 3, 2)),
@@ -58,6 +60,7 @@ def test_weightnorm_dense_train_notinit(self):
         self.assertTrue(hasattr(model.layers[0], 'g'))
 
     def test_weightnorm_conv2d(self):
+        # TODO: Get data init to work with tf_function compile #428
         model = tf.keras.models.Sequential()
         model.add(
             wrappers.WeightNormalization(
@@ -67,7 +70,8 @@ def test_weightnorm_conv2d(self):
         model.add(tf.keras.layers.Activation('relu'))
         model.compile(
             optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.001),
-            loss='mse')
+            loss='mse',
+            experimental_run_tf_function=False)
         model.fit(
             np.random.random((2, 4, 4, 3)),
             np.random.random((2, 4, 4, 5)),