Modified InputSpec of SpectralNormalization layer (#21335)

* Modified InputSpec of SpectralNormalization layer

* Added unit test for for higher dimension

Author: sonali-kumari1

keras/src/layers/normalization/spectral_normalization.py

@@ -52,7 +52,7 @@ def __init__(self, layer, power_iterations=1, **kwargs):
 
     def build(self, input_shape):
         super().build(input_shape)
-        self.input_spec = InputSpec(shape=[None] + list(input_shape[1:]))
+        self.input_spec = InputSpec(min_ndim=1, axes={-1: input_shape[-1]})
 
         if hasattr(self.layer, "kernel"):
             self.kernel = self.layer.kernel

keras/src/layers/normalization/spectral_normalization_test.py

@@ -35,6 +35,20 @@ def test_basic_spectralnorm(self):
             run_training_check=False,
         )
 
+    @pytest.mark.requires_trainable_backend
+    def test_spectralnorm_higher_dim(self):
+        self.run_layer_test(
+            layers.SpectralNormalization,
+            init_kwargs={"layer": layers.Dense(2)},
+            input_data=np.random.uniform(size=(10, 3, 4, 5)),
+            expected_output_shape=(10, 3, 4, 2),
+            expected_num_trainable_weights=2,
+            expected_num_non_trainable_weights=1,
+            expected_num_seed_generators=0,
+            expected_num_losses=0,
+            supports_masking=False,
+        )
+
     def test_invalid_power_iterations(self):
         with self.assertRaisesRegex(
             ValueError, "`power_iterations` should be greater than zero."