[Bug] Gemma3nAudioConverter fails with NotImplementedError for tf.Tensor inputs under JAX/PyTorch backends

[maitry63]

Gemma3nAudioConverter fails with symbolic tf.Tensor under JAX and PyTorch backends

Description

While working on PR #2978 for the PyGrain migration of Gemma3n, CI failures were observed in Gemma3nCausalLMTest under the JAX and PyTorch backends.

Affected tests:

• test_causal_lm_basics_audio_text
• test_causal_lm_basics_multimodal

The failures occur when Gemma3nAudioConverter is called from the Gemma3n preprocessing pipeline with a symbolic tf.Tensor.

KerasHub can use JAX or PyTorch as the active Keras backend while TensorFlow is still used for TensorFlow-based preprocessing/data-pipeline execution. Therefore, a TensorFlow tensor can legitimately reach a preprocessing layer even when the active Keras backend is JAX or PyTorch.

Currently, Gemma3nAudioConverter does not correctly handle this TensorFlow/non-TensorFlow backend boundary.

[Note: This issue may also affect other audio converter models when their preprocessing pipelines are migrated to PyGrain, since they can encounter the same TensorFlow tensor → non-TensorFlow backend boundary during preprocessing.]

---

Reproducable colab.

Minimal Reproduction

import os

# Set this to "jax" or "torch" to reproduce the issue.
os.environ["KERAS_BACKEND"] = "jax"

import keras
import tensorflow as tf

from keras_hub.src.models.gemma3n.gemma3n_audio_converter import (
    Gemma3nAudioConverter,
)

print(f"Active Keras backend: {keras.config.backend()}")

converter = Gemma3nAudioConverter(
    feature_size=16,
    sampling_rate=16000,
    padding_value=0.0,
    return_attention_mask=True,
    frame_length_ms=32.0,
    hop_length_ms=10.0,
    min_frequency=125.0,
    max_frequency=7600.0,
    preemphasis=0.97,
    preemphasis_htk_flavor=True,
    fft_overdrive=True,
    dither=0.0,
    input_scale_factor=1.0,
    mel_floor=1e-5,
    per_bin_mean=None,
    per_bin_stddev=None,
    padding_side="right",
)


@tf.function
def preprocess_step(raw_audio):
    return converter(raw_audio)


audio_input = tf.ones(
    (1, 16000),
    dtype=tf.float32,
)

print("\nRunning converter with a symbolic tf.Tensor...")

try:
    outputs = preprocess_step(audio_input)
    print("Success:", outputs)
except Exception as e:
    print(f"\n{type(e).__name__}: {e}")

Output:

Active Keras backend: jax

Running converter with a symbolic tf.Tensor...

NotImplementedError: Cannot convert a symbolic tf.Tensor to a numpy array.

Cannot convert a symbolic tf.Tensor (raw_audio:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported.
    
    Arguments received by Gemma3nAudioConverter.call():
      • raw_speech=<tf.Tensor 'raw_audio:0' shape=(1, 16000) dtype=float32>
      • padding=longest
      • max_length=480000
      • truncation=True
      • pad_to_multiple_of=128
      • return_attention_mask=True

Root Cause:

• Gemma3nAudioConverter can receive a symbolic tf.Tensor from the TensorFlow preprocessing/data pipeline.
• The active Keras backend can be JAX or PyTorch.
• ops.is_tensor() is backend-specific, so it does not identify the TensorFlow tensor under JAX/PyTorch.
• The TensorFlow tensor therefore reaches the NumPy conversion path.
• A symbolic tf.Tensor cannot be converted to NumPy inside tf.function, causing the failure.

Expected Behavior:

Gemma3nAudioConverter should correctly handle a symbolic tf.Tensor received from the TensorFlow preprocessing/data pipeline, even when the active Keras backend is JAX or PyTorch.

The converter should process the input without attempting to convert the symbolic TensorFlow tensor to NumPy during graph tracing.

This issue was discovered while working on:
#2978 — PyGrain migration of Gemma3n.

Working on this fix.