Spectral Normalization layer does not work

[dryglicki]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 22.04.4 LTS
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.18.0
• Python version: 3.11.11
• Bazel version (if compiling from source): N/A
• GPU model and memory: RTX Ada A6000, 48 GB VRAM

Describe the problem.

The SpectralNormalization Wrapper does not work.

Describe the current behavior.

When used, it produces the following error:

Traceback (most recent call last):
  File "/home/dryglicki/code/testing/test_specnorm.py", line 15, in <module>
    x = KL.SpectralNormalization(KL.Conv2D(32, 3, padding='same'))(inputs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/engine/base_layer.py", line 1288, in input_spec
    raise TypeError(
TypeError: Layer input_spec must be an instance of InputSpec. Got: InputSpec(shape=(None, None, None, 12), ndim=4)

Describe the expected behavior.

For SpectralNormalization to work.

Contributing.

• Do you want to contribute a PR? (yes/no): no

Standalone code to reproduce the issue.

import tensorflow as tf
import tf_keras as K
import tf_keras.layers as KL

nb = 15
nt = 6
nx = 32
ny = 32
nc = 12

xx = tf.random.normal((nb, nx, ny, nc))
yy = tf.random.normal((nb, nx, ny, nc))

inputs = KL.Input( [None, None, nc], name = 'dummy_input')
x = KL.SpectralNormalization(KL.Conv2D(32, 3, padding='same'))(inputs)
# x = KL.Conv2D(32, 3, padding='same')(inputs)
x = KL.Conv2D(nc, 1, padding='same')(x)

model = K.models.Model(inputs = inputs, outputs = x)

model.compile(loss = 'mae')

model.fit(xx,yy,epochs = 10)

Swapping the Conv2D(32) lines results in functioning code.

[dryglicki]

Ah ha, found it:

https://github.com/keras-team/tf-keras/blob/master/tf_keras/layers/normalization/spectral_normalization.py#L72-L74

It's a mismatch between tf_keras and tf.keras. Sure, I could issue a legacy call to Keras, but this is the heart of the problem.

[dryglicki]

More fun!

If I make the switch about, switching tf.keras for tf_keras at the InputSpec line, I am then greeted by this error:

Traceback (most recent call last):
  File "/home/dryglicki/code/testing/test_specnorm.py", line 161, in <module>
    model.fit(xx,yy,epochs = 10)
  File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/autograph_util.py", line 52, in autograph_handler
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: in user code:

    File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/engine/training.py", line 1398, in train_function  *
        return step_function(self, iterator)
    File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/engine/training.py", line 1381, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/engine/training.py", line 1370, in run_step  **
        outputs = model.train_step(data)
    File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/engine/training.py", line 1147, in train_step
        y_pred = self(x, training=True)
    File "/ssd0/miniforge3_2024-04/envs/tensorflow_2d18_py3d11/lib/python3.11/site-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None

    OperatorNotAllowedInGraphError: Exception encountered when calling layer 'spectral_normalization' (type SpectralNormalization).
    
    Using a symbolic `tf.Tensor` as a Python `bool` is not allowed. You can attempt the following resolutions to the problem: If you are running in Graph mode, use Eager execution mode or decorate this function with @tf.function. If you are using AutoGraph, you can try decorating this function with @tf.function. If that does not work, then you may be using an unsupported feature or your source code may not be visible to AutoGraph. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md#access-to-source-code for more information.
    
    Call arguments received by layer 'spectral_normalization' (type SpectralNormalization):
      • inputs=tf.Tensor(shape=(None, 32, 32, 12), dtype=float32)
      • training=True

The offending line in Spectral Normalization then becomes: https://github.com/keras-team/tf-keras/blob/master/tf_keras/layers/normalization/spectral_normalization.py#L119

Commenting that out and de-indenting as appropriate finally makes SpectralNormalization function properly.