Fixes handling of input tensor rank.

PiperOrigin-RevId: 360561777
Change-Id: I92f15753c033c40d1965c8d9e1e3ce6fa0f25848

Author: Johannes Ballé

tensorflow_compression/python/layers/gdn.py

@@ -152,7 +152,6 @@ def __init__(self,
       **kwargs: Other keyword arguments passed to superclass (`Layer`).
     """
     super().__init__(**kwargs)
-    self.input_spec = tf.keras.layers.InputSpec(min_ndim=2)
     self.inverse = inverse
     self.rectify = rectify
     self.data_format = data_format
@@ -329,14 +328,13 @@ def _channel_axis(self):
     return {"channels_first": 1, "channels_last": -1}[self.data_format]
 
   def build(self, input_shape):
-    channel_axis = self._channel_axis
     input_shape = tf.TensorShape(input_shape)
-    num_channels = input_shape[channel_axis]
+    if input_shape.rank is None or input_shape.rank < 2:
+      raise ValueError(f"Input tensor must have at least rank 2, received "
+                       f"shape {input_shape}.")
+    num_channels = input_shape[self._channel_axis]
     if num_channels is None:
-      raise ValueError("The channel dimension of the inputs to `GDN` "
-                       "must be defined.")
-    self.input_spec = tf.keras.layers.InputSpec(
-        min_ndim=2, axes={channel_axis: num_channels})
+      raise ValueError("The channel dimension of the inputs must be defined.")
 
     if self.alpha_parameter is None:
       initial_value = self.alpha_initializer(
@@ -367,8 +365,9 @@ def build(self, input_shape):
   def call(self, inputs) -> tf.Tensor:
     inputs = tf.convert_to_tensor(inputs, dtype=self.dtype)
     rank = inputs.shape.rank
-    if rank is None:
-      raise RuntimeError("Input tensor rank must be defined.")
+    if rank is None or rank < 2:
+      raise ValueError(f"Input tensor must have at least rank 2, received "
+                       f"shape {inputs.shape}.")
 
     if self.rectify:
       inputs = tf.nn.relu(inputs)

tensorflow_compression/python/layers/signal_conv.py

@@ -315,7 +315,6 @@ def __init__(self, filters, kernel_support,
       **kwargs: Keyword arguments passed to superclass (`Layer`).
     """
     super().__init__(**kwargs)
-    self.input_spec = tf.keras.layers.InputSpec(ndim=self._rank + 2)
     self.filters = filters
     self.kernel_support = kernel_support
     self.corr = corr
@@ -574,12 +573,13 @@ def _raise_notimplemented(self):
 
   def build(self, input_shape):
     input_shape = tf.TensorShape(input_shape)
+    if input_shape.rank != self._rank + 2:
+      raise ValueError(f"Input tensor must have rank {self._rank + 2}, "
+                       f"received shape {input_shape}.")
     channel_axis = {"channels_first": 1, "channels_last": -1}[self.data_format]
     input_channels = input_shape[channel_axis]
     if input_channels is None:
       raise ValueError("The channel dimension of the inputs must be defined.")
-    self.input_spec = tf.keras.layers.InputSpec(
-        ndim=self._rank + 2, axes={channel_axis: input_channels})
 
     kernel_shape = self.kernel_support + (input_channels, self.filters)
     if self.channel_separable:
@@ -837,7 +837,10 @@ def _up_convolve_transpose_explicit(self, inputs, kernel, prepadding):
     return outputs
 
   def call(self, inputs) -> tf.Tensor:
-    inputs = tf.convert_to_tensor(inputs)
+    inputs = tf.convert_to_tensor(inputs, dtype=self.dtype)
+    if inputs.shape.rank != self._rank + 2:
+      raise ValueError(f"Input tensor must have rank {self._rank + 2}, "
+                       f"received shape {inputs.shape}.")
     outputs = inputs
 
     # Not for all possible combinations of (`kernel_support`, `corr`,

tensorflow_compression/python/layers/signal_conv_test.py

@@ -31,15 +31,15 @@ def test_invalid_data_format_raises_error(self):
 
   def test_variables_are_enumerated(self):
     layer = signal_conv.SignalConv2D(3, 1, use_bias=True)
-    layer.build((None, None, 2))
+    layer.build((None, None, None, 2))
     self.assertLen(layer.weights, 2)
     self.assertLen(layer.trainable_weights, 2)
     weight_names = [w.name for w in layer.weights]
     self.assertSameElements(weight_names, ["kernel_rdft:0", "bias:0"])
 
   def test_bias_variable_is_not_unnecessarily_created(self):
     layer = signal_conv.SignalConv2D(5, 3, use_bias=False)
-    layer.build((None, None, 3))
+    layer.build((None, None, None, 3))
     self.assertLen(layer.weights, 1)
     self.assertLen(layer.trainable_weights, 1)
     weight_names = [w.name for w in layer.weights]
@@ -49,14 +49,14 @@ def test_variables_are_not_enumerated_when_overridden(self):
     layer = signal_conv.SignalConv2D(1, 1)
     layer.kernel_parameter = [[[[1]]]]
     layer.bias_parameter = [0]
-    layer.build((None, 1))
+    layer.build((None, None, None, 1))
     self.assertEmpty(layer.weights)
     self.assertEmpty(layer.trainable_weights)
 
   def test_variables_trainable_state_follows_layer(self):
     layer = signal_conv.SignalConv2D(1, 1, use_bias=True)
     layer.trainable = False
-    layer.build((None, 1))
+    layer.build((None, None, None, 1))
     self.assertLen(layer.weights, 2)
     self.assertEmpty(layer.trainable_weights)