Add audio support to Gemma4 sampler.

tests/models/gemma4/audio_test.py

@@ -0,0 +1,55 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for audio."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from tunix.models.gemma4 import audio
+
+
+class AudioTest(parameterized.TestCase):
+
+  def test_audio_encoder_shape(self):
+    config = audio.AudioEncoderConfig(
+        num_layers=2,  # fewer layers for faster test
+        model_dims=128,
+        lm_model_dims=256,
+        atten_num_heads=4,
+    )
+
+    rngs = nnx.Rngs(0)
+    model = audio.AudioEncoder(config, rngs=rngs)
+
+    batch_size = 2
+    num_samples = 16000  # 1 second
+
+    x = jax.random.normal(jax.random.PRNGKey(1), (batch_size, num_samples))
+    sequence_lengths = jnp.full((batch_size,), num_samples, dtype=jnp.int32)
+
+    output, mask = model(x, sequence_lengths)
+
+    expected_seq_len = 25
+    expected_shape = (batch_size, expected_seq_len, config.lm_model_dims)
+
+    self.assertEqual(output.shape, expected_shape)
+    self.assertEqual(mask.shape, (batch_size, expected_seq_len))
+    self.assertFalse(jnp.any(jnp.isnan(output)))
+
+
+if __name__ == "__main__":
+  absltest.main()

tests/models/gemma4/model_test.py

@@ -220,7 +220,7 @@ def test_forward_pass_vision(self):
     tokens = jax.random.randint(
         jax.random.PRNGKey(0), (1, 32), 0, config.num_embed
     )
-    tokens = tokens.at[0, 10:15].set(model_lib.TOKEN_PLACEHOLDER)
+    tokens = tokens.at[0, 10:15].set(model_lib.VISION_TOKEN_PLACEHOLDER)
 
     positions = jnp.tile(
         jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1)
@@ -277,7 +277,7 @@ def test_forward_pass_vision_bidirectional(self):
     tokens = jax.random.randint(
         jax.random.PRNGKey(0), (1, 32), 0, config.num_embed
     )
-    tokens = tokens.at[0, 10:15].set(model_lib.TOKEN_PLACEHOLDER)
+    tokens = tokens.at[0, 10:15].set(model_lib.VISION_TOKEN_PLACEHOLDER)
 
     positions = jnp.tile(
         jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1)
@@ -337,9 +337,9 @@ def test_forward_pass_vision_batch(self):
         jax.random.PRNGKey(0), (batch_size, seq_len), 0, config.num_embed
     )
     # Image placeholders: token shape represents visual soft tokens within sequences.
-    tokens = tokens.at[0, 10:15].set(model_lib.TOKEN_PLACEHOLDER)
-    tokens = tokens.at[1, 5:8].set(model_lib.TOKEN_PLACEHOLDER)
-    tokens = tokens.at[1, 20:25].set(model_lib.TOKEN_PLACEHOLDER)
+    tokens = tokens.at[0, 10:15].set(model_lib.VISION_TOKEN_PLACEHOLDER)
+    tokens = tokens.at[1, 5:8].set(model_lib.VISION_TOKEN_PLACEHOLDER)
+    tokens = tokens.at[1, 20:25].set(model_lib.VISION_TOKEN_PLACEHOLDER)
 
     positions = jnp.tile(
         jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1)
@@ -377,6 +377,117 @@ def test_forward_pass_vision_batch(self):
     )
     self.assertEqual(logits.shape, (batch_size, seq_len, config.num_embed))
 
+  def test_forward_pass_audio(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.num_layers = 1
+    config.embed_dim = 256
+    config.hidden_dim = 512
+    config.num_heads = 4
+    config.head_dim = 64
+    config.num_kv_heads = 1
+    config.frac_shared_layers = 0.0
+    config.audio_encoder = model_lib.audio.AudioEncoderConfig(
+        num_layers=1,
+        model_dims=64,
+        lm_model_dims=128,
+        atten_num_heads=2,
+    )
+
+    rngs = nnx.Rngs(0)
+    model = model_lib.Gemma4(config, rngs=rngs, text_only=False)
+
+    tokens = jax.random.randint(
+        jax.random.PRNGKey(0), (1, 32), 0, config.num_embed
+    )
+    tokens = tokens.at[0, 10:15].set(model_lib.AUDIO_TOKEN_PLACEHOLDER)
+
+    positions = jnp.tile(
+        jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1)
+    )
+    attn_mask = jnp.tril(
+        jnp.ones((tokens.shape[1], tokens.shape[1]), dtype=jnp.bool_)
+    )[None, ...]
+
+    soft_token_counts = (5,)
+    audio_data = jnp.zeros((1, 1, 16000), dtype=jnp.float32)
+    audio_lengths = jnp.full((1, 1), 16000, dtype=jnp.int32)
+
+    audios = model_lib.PreprocessedAudioInput(
+        audio=audio_data,
+        audio_lengths=audio_lengths,
+        soft_token_counts=soft_token_counts,
+    )
+
+    logits, _ = model(
+        tokens,
+        positions=positions,
+        attention_mask=attn_mask,
+        audios=audios,
+    )
+    self.assertEqual(logits.shape, (1, 32, config.num_embed))
+
+  def test_forward_pass_audio_batch(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.num_layers = 1
+    config.embed_dim = 256
+    config.hidden_dim = 512
+    config.num_heads = 4
+    config.head_dim = 64
+    config.num_kv_heads = 1
+    config.frac_shared_layers = 0.0
+    config.audio_encoder = model_lib.audio.AudioEncoderConfig(
+        num_layers=1,
+        model_dims=64,
+        lm_model_dims=128,
+        atten_num_heads=2,
+    )
+
+    rngs = nnx.Rngs(0)
+    model = model_lib.Gemma4(config, rngs=rngs, text_only=False)
+
+    batch_size = 2
+    seq_len = 32
+    tokens = jax.random.randint(
+        jax.random.PRNGKey(0), (batch_size, seq_len), 0, config.num_embed
+    )
+    tokens = tokens.at[0, 10:15].set(model_lib.AUDIO_TOKEN_PLACEHOLDER)
+    tokens = tokens.at[1, 5:8].set(model_lib.AUDIO_TOKEN_PLACEHOLDER)
+    tokens = tokens.at[1, 20:25].set(model_lib.AUDIO_TOKEN_PLACEHOLDER)
+
+    positions = jnp.tile(
+        jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1)
+    )
+    attn_mask = jnp.tril(
+        jnp.ones((tokens.shape[1], tokens.shape[1]), dtype=jnp.bool_)
+    )[None, ...]
+    attn_mask = jnp.broadcast_to(attn_mask, (batch_size, seq_len, seq_len))
+
+    soft_token_counts = ((5,), (3, 5))
+    max_n_audios = 2
+    max_samples = 16000
+
+    audio_data = jnp.zeros(
+        (batch_size, max_n_audios, max_samples), dtype=jnp.float32
+    )
+    audio_lengths = jnp.full(
+        (batch_size, max_n_audios), max_samples, dtype=jnp.int32
+    )
+    audio_lengths = audio_lengths.at[0, 1].set(0)
+
+    audios = model_lib.PreprocessedAudioInput(
+        audio=audio_data,
+        audio_lengths=audio_lengths,
+        soft_token_counts=soft_token_counts,
+    )
+
+    logits, _ = model(
+        tokens,
+        positions=positions,
+        attention_mask=attn_mask,
+        audios=audios,
+    )
+    self.assertEqual(logits.shape, (batch_size, seq_len, config.num_embed))
+
 
 if __name__ == "__main__":
   absltest.main()

tests/processors/audio_processor_test.py

@@ -0,0 +1,98 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for audio_processor."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import numpy as np
+from tunix.processors import audio_processor
+
+
+class AudioProcessorTest(parameterized.TestCase):
+
+  def test_compute_soft_token_count(self):
+    # 1 second of audio at 16kHz
+    self.assertEqual(audio_processor.compute_soft_token_count(16000), 25)
+    # 2 seconds of audio at 16kHz
+    self.assertEqual(audio_processor.compute_soft_token_count(32000), 50)
+
+  def test_add_variable_extra_tokens_for_audio(self):
+    tokens = np.array([
+        [1, 2, 258881, 3],
+        [4, 258881, 258881, 5],
+    ])
+    soft_token_counts = ((2,), (1, 2))
+
+    expanded = audio_processor.add_variable_extra_tokens_for_audio(
+        tokens,
+        soft_token_counts=soft_token_counts,
+    )
+
+    expected_row_0 = [1, 2, 256000, -4, -4, 258883, 3]
+    expected_row_1 = [4, 256000, -4, 258883, 256000, -4, -4, 258883, 5]
+
+    max_len = max(len(expected_row_0), len(expected_row_1))
+    padded_expected = np.zeros((2, max_len), dtype=np.int32)
+    padded_expected[0, : len(expected_row_0)] = expected_row_0
+    padded_expected[1, : len(expected_row_1)] = expected_row_1
+
+    np.testing.assert_array_equal(expanded, padded_expected)
+
+  def test_process_gemma4_audio_inputs_batch(self):
+    class _DummyAudioConfig:
+      sample_rate = 16000
+      audio_seq_length = 750
+
+    config = _DummyAudioConfig()
+
+    audio = [[np.zeros(16000), np.zeros(8000)], [np.zeros(32000)]]
+    tokens = [
+        np.array([1, 258881, 3, 258881]),
+        np.array([258881, 4]),
+    ]
+
+    processed, new_tokens = audio_processor.process_gemma4_audio_inputs(
+        audio=audio,
+        tokens=tokens,
+        audio_encoder_config=config,
+        pad_id=0,
+    )
+
+    self.assertEqual(processed.audio.shape, (2, 2, 32000))
+    self.assertEqual(processed.audio_lengths.shape, (2, 2))
+
+    np.testing.assert_array_equal(processed.audio_lengths[0], [16000, 8000])
+    np.testing.assert_array_equal(processed.audio_lengths[1], [32000, 0])
+
+    self.assertEqual(processed.soft_token_counts, ((25, 12), (50,)))
+
+    self.assertLen(new_tokens, 2)
+    expected_new_tokens_0 = np.array(
+        [1]
+        + [256000]
+        + [-4] * 25
+        + [258883]
+        + [3]
+        + [256000]
+        + [-4] * 12
+        + [258883]
+    )
+    expected_new_tokens_1 = np.array([256000] + [-4] * 50 + [258883] + [4])
+    np.testing.assert_array_equal(new_tokens[0], expected_new_tokens_0)
+    np.testing.assert_array_equal(new_tokens[1], expected_new_tokens_1)
+
+
+if __name__ == "__main__":
+  absltest.main()

tunix/generate/sampler.py

@@ -36,6 +36,7 @@
 from tunix.generate import utils
 import tunix.generate.beam_search as beam_search_lib
 import tunix.generate.tokenizer_adapter as tok_adapter
+from tunix.processors import audio_processor
 from tunix.processors import image_processor
 
 LayerCache = dict[str, jaxtyping.Array]
@@ -562,6 +563,7 @@ def _prefill_fn(
       params: statelib.State,
       sampler_state: _SamplingState,
       images: jnp.ndarray | None = None,
+      audios: Any = None,
       echo: bool = True,
   ) -> _SamplingState:
     """Performs prefill."""
@@ -600,7 +602,11 @@ def _prefill_fn(
       )
 
     transformer = nnx.merge(self._transformer_graphdef, params)
-    kwargs = {} if images is None else {'images': images}
+    kwargs = {}
+    if images is not None:
+      kwargs['images'] = images
+    if audios is not None:
+      kwargs['audios'] = audios
     decode_only_last_token = self._supports_decode_only_last_token and not echo
     if decode_only_last_token:
       kwargs['decode_only_last_token'] = True
@@ -759,6 +765,13 @@ def __call__(
           | jnp.ndarray
           | None
       ) = None,
+      audios: (
+          str
+          | np.ndarray
+          | list[str | np.ndarray | list[str | np.ndarray] | None]
+          | jnp.ndarray
+          | None
+      ) = None,
   ) -> base_sampler.SamplerOutput:
     """Samples a completion of the input string.
 
@@ -804,16 +817,16 @@ def __call__(
     tokens = [self.tokenize(x) for x in input_strings]
 
     processed_images = images
-    is_gemma4_multimodal = (
+    is_gemma4_vision = (
         hasattr(self.transformer, 'vision_encoder')
         and self.transformer.vision_encoder is not None
     )
 
-    if is_gemma4_multimodal and images is not None:
+    if is_gemma4_vision and images is not None:
       processed_images, tokens = image_processor.process_gemma4_inputs(
           images,
           tokens,
-          self.transformer.vision_encoder,
+          getattr(self.transformer, 'vision_encoder'),
           self.tokenizer.pad_id(),
       )