Problem

The Model.export() API in Keras 3 supports exporting to a TensorFlow SavedModel artifact for inference. When trying to export Gemma 2 and ShieldGemma to TF SavedModel, I ran into two different ValueErrors:

• If no input_signature is provided, a ValueError will be thrown related to a structural mismatch between the expected and actual inputs passed to the GemmaCausalLM class; and
• If an input_signature is provided as a list[keras.InputSpec], a ValueError will be thrown related the the wrong number of values being passed to a TF function.

However, if yo uwrap the dict from model.input in a list, as input_signature=[model.input], the export runs to completion.

This is not restricted to Gemma models, as shown in this minimal reproducible example.

Thanks to @mattdangerw for helping to isolate this minimal example.