Keras 3 SpectralNormalization() issue

[Kwokleung Chan]

There seems a bug with Keras 3  SpectralNormalization

Following https://www.tensorflow.org/api_docs/python/tf/keras/layers/SpectralNormalization

The below works in keras 3.0.5

import keras

from keras import layers
import numpy as np

x = np.random.rand(1, 10, 10, 1)
dense = layers.SpectralNormalization(keras.layers.Dense(10))
y = dense(x)
y.shape

Out[1]: TensorShape([1, 10, 10, 10])

but the below will have error.

import keras
from keras import layers
import numpy as np

# Prepare the dataset. We use both the training & test MNIST digits.
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
num_classes = 10

x_train = x_train.astype("float32") / 255.0
y_train = keras.utils.to_categorical(y_train, num_classes)

Input = keras.Input(shape=(28,28,1))
x = keras.layers.Flatten()(Input)
dense = layers.SpectralNormalization(keras.layers.Dense(num_classes, activation = 'softmax'))
Output = dense(x)

model = keras.Model(inputs = Input, outputs = Output)
model.compile(loss = keras.losses.CategoricalCrossentropy(), metrics = ['accuracy'], optimizer = 'Adam')
model.summary()

model.fit(x_train, y_train, epochs = 2)

OperatorNotAllowedInGraphError: Exception encountered when calling SpectralNormalization.call().

Iterating over a symbolic `tf.Tensor` 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.

Arguments received by SpectralNormalization.call():
  • inputs=tf.Tensor(shape=(32, 784), dtype=float32)
  • training=True