Don't overspecify WeightNorm input_spec, match input_spec of wrapped

PiperOrigin-RevId: 359573317

Author: T2T Team

tensor2tensor/layers/common_layers.py

@@ -4043,9 +4043,6 @@ def _data_dep_init(self, inputs):
 
   def build(self, input_shape=None):
     """Build `Layer`."""
-    input_shape = tf.TensorShape(input_shape).as_list()
-    self.input_spec = layers().InputSpec(shape=input_shape)
-
     if not self.layer.built:
       self.layer.build(input_shape)
       self.layer.built = False
@@ -4072,6 +4069,7 @@ def build(self, input_shape=None):
       self._compute_weights()
 
       self.layer.built = True
+    self.input_spec = self.layer.input_spec
 
     super(WeightNorm, self).build()
     self.built = True

tensor2tensor/layers/common_layers_test.py

@@ -965,5 +965,20 @@ def fn_recompute(x):
         self.assertAllClose(g1, g2)
 
 
+class WeightNormTest(tf.test.TestCase):
+
+  def testInputSpec(self):
+    """Test that WeighNorm does not overspecify the input_spec."""
+    conv = common_layers.WeightNorm(
+        tf.keras.layers.Conv1D(filters=8, kernel_size=3))
+    # Call with one batch size:
+    conv(tf.zeros([1, 16, 2]))
+    # Should allow call with another batch size.
+    conv(tf.zeros([2, 16, 2]))
+    # Input spec does detect incorrect input feature dim.
+    with self.assertRaises(ValueError):
+      conv(tf.zeros([2, 16, 3]))
+
+
 if __name__ == "__main__":
   tf.test.main()