Migrate: Gemma3nAudioConverter from TF ops to keras.ops

[maitry63]

This PR migrates Gemma3nAudioConverter from TensorFlow ops to keras.ops.

Replaces TensorFlow-specific audio operations with Keras equivalents, including

• tf.signal.frame → keras.ops.extract_sequences
• tf.signal.rfft → keras.ops.rfft,
  while preserving batched/unbatched audio handling, spectrogram padding, and attention-mask behavior.

Fixes: #2951 (1st part)

Numerical verification colab.

Numerical Verification Results

Verified numerical equivalence against pure NumPy reference :

• Test 1 (Sine Wave): Log-Mel max diff = 9.64e-03 | Linear-Mel max diff = 1.14e-04 [PASS]
• Test 2 (Normalized Audio): max diff = 2.19e-05 [PASS]
• Test 3 (Silent Audio Floor): max diff = 3.17e-07 [PASS]
• Test 4 (Batch Invariance): max diff = 0.00e+00 [PASS]

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have followed the Keras Hub Model contribution guidelines in making these changes.
• I have followed the Keras Hub API design guidelines in making these changes.
• I have signed the Contributor License Agreement.

[gemini-code-assist]

Code Review

This pull request refactors the Gemma3nAudioConverter to be backend-agnostic by replacing TensorFlow-specific operations with Keras ops and NumPy equivalents, and adds corresponding unit tests. While the migration is a great step, three critical issues were identified in the review: a runtime unpacking error due to ops.rfft returning a complex tensor instead of a tuple, a potential crash during symbolic tracing when calling np.asarray on a symbolic KerasTensor, and performance/compatibility bottlenecks on GPUs and TPUs caused by casting computations to float64 instead of using self.compute_dtype.

[gemini-code-assist]

In Keras 3, ops.rfft returns a complex-valued tensor rather than a tuple of (real, imag). Unpacking it directly will result in a runtime ValueError (e.g., not enough values to unpack). To compute the magnitude of the STFT, you should use ops.abs on the complex tensor returned by ops.rfft.

Suggested change

	real, imag = ops.rfft(fft_frames, fft_length=self.fft_length)
	magnitude_spec = ops.sqrt(ops.square(real) + ops.square(imag))
	stft = ops.rfft(fft_frames, fft_length=self.fft_length)
	magnitude_spec = ops.abs(stft)

[gemini-code-assist]

The layer will crash when called with a symbolic KerasTensor (e.g., during functional model building or tracing) because np.asarray(raw_speech) is called directly on the input. To support symbolic tracing and functional model building, we should add a defensive check at the beginning of call() to return symbolic KerasTensors with the expected shapes and dtypes when a symbolic tensor is passed.

        import keras
        if isinstance(raw_speech, keras.KerasTensor):
            was_batched = len(raw_speech.shape) > 1
            if was_batched:
                features_shape = (raw_speech.shape[0], None, self.feature_size)
                mask_shape = (raw_speech.shape[0], None)
            else:
                features_shape = (None, self.feature_size)
                mask_shape = (None,)
            return (
                keras.KerasTensor(shape=features_shape, dtype=self.compute_dtype),
                keras.KerasTensor(shape=mask_shape, dtype="int32"),
            )

        if isinstance(raw_speech, (list, tuple)):
            was_batched = True
            raw_speech_list = [np.asarray(speech) for speech in raw_speech]

References

1. When handling inputs that may be passed as tensors in a Python-only execution path, perform defensive type conversion or handling at the entry point of the path rather than deep inside helper functions.

[gemini-code-assist]

Casting to float64 for the entire spectrogram extraction in call() is a major performance bottleneck on GPUs and is not supported on TPUs (which can cause compilation errors or silent CPU fallbacks). Since this is a preprocessing layer, using self.compute_dtype (which defaults to float32) is much more efficient and device-compatible. Please update this and the other float64 casts in this method to use self.compute_dtype.

Suggested change

	fft_frames = ops.cast(frames, dtype="float64")
	fft_frames = ops.cast(frames, dtype=self.compute_dtype)