Remove compute_output_shape.

Keras: "manual" shape inference is only required if the layer is dynamic (otherwise we use TF's static shape inference capabilities)

PiperOrigin-RevId: 290821518

official/nlp/modeling/layers/attention.py

@@ -118,14 +118,6 @@ class Attention(tf.keras.layers.Layer):
 
     self._dropout = tf.keras.layers.Dropout(rate=self._dropout_rate)
 
-  def compute_output_shape(self, input_shape):
-    # TODO(momernick): validate tensor dimensions.
-    from_tensor_shape = tf.TensorShape(input_shape[0])
-    batch = from_tensor_shape[0]
-    from_tensor_length = from_tensor_shape[1]
-    return tf.TensorShape(
-        (batch, from_tensor_length, self._num_heads, self._head_size))
-
   def get_config(self):
     config = {
         "num_heads":

official/nlp/modeling/layers/dense_einsum.py

@@ -143,18 +143,6 @@ class DenseEinsum(tf.keras.layers.Layer):
       self._bias = None
     super(DenseEinsum, self).build(input_shape)
 
-  def compute_output_shape(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    input_shape = input_shape.with_rank_at_least(self._num_summed_dimensions +
-                                                 1)
-    for i in range(self._num_summed_dimensions):
-      if tf.dimension_value(input_shape[-1 * i]) is None:
-        raise ValueError(
-            "The %s dimension of input_shape must be defined, but saw: %s" %
-            (-1 * i, input_shape))
-    return input_shape[:-1 * self._num_summed_dimensions].concatenate(
-        self._units)
-
   def get_config(self):
     config = {
         "output_shape":

official/nlp/modeling/layers/transformer.py

@@ -158,13 +158,6 @@ class Transformer(tf.keras.layers.Layer):
 
     super(Transformer, self).build(input_shape)
 
-  def compute_output_shape(self, input_shape):
-    data_tensor_shape = tf.TensorShape(input_shape[0])
-    batch = data_tensor_shape[0]
-    sequence_length = data_tensor_shape[1]
-
-    return tf.TensorShape((batch, sequence_length, self._output_einsum_shape))
-
   def get_config(self):
     config = {
         "num_attention_heads":

official/nlp/modeling/layers/transformer_scaffold.py

@@ -175,13 +175,6 @@ class TransformerScaffold(tf.keras.layers.Layer):
 
     super(TransformerScaffold, self).build(input_shape)
 
-  def compute_output_shape(self, input_shape):
-    data_tensor_shape = tf.TensorShape(input_shape[0])
-    batch = data_tensor_shape[0]
-    sequence_length = data_tensor_shape[1]
-
-    return tf.TensorShape((batch, sequence_length, self._output_einsum_shape))
-
   def get_config(self):
     config = {
         "attention_cls":

official/nlp/modeling/networks/masked_lm.py

@@ -168,9 +168,6 @@ class Bias(tf.keras.layers.Layer):
 
     super(Bias, self).build(input_shape)
 
-  def compute_output_shape(self, input_shape):
-    return input_shape
-
   def get_config(self):
     config = {
         'activation': tf.keras.activations.serialize(self._activation),