optional PLE token input

PiperOrigin-RevId: 906871699

Author: The gemma Authors

gemma/gm/nn/gemma4/_modules.py

@@ -166,13 +166,20 @@ def encode_audio(self, x: jax.Array) -> jax.Array:
     x = self.audio_soft_embedding_norm(x)
     return x
 
-  def encode_per_layer_input(self, x: jax.Array, t: jax.Array) -> jax.Array:
+  def encode_per_layer_input(
+      self,
+      x: jax.Array,
+      t: jax.Array,
+      ignore_ple_tokens: bool = False,
+  ) -> jax.Array:
     """Encodes the input tokens.
 
     Args:
       x: Input shape [seq_len, embed_dim] or [batch_size, seq_len, embed_dim].
       t: Input tokens of shape [seq_len] or [batch_size, seq_len], where each
         token is an integer in [0, vocab_size).
+      ignore_ple_tokens: If True, the tokens are not used to compute the per
+        layer input embeddings.
 
     Returns:
       Encoded input of shape [seq_len, num_layers, per_layer_input_dim] or
@@ -184,6 +191,8 @@ def encode_per_layer_input(self, x: jax.Array, t: jax.Array) -> jax.Array:
     )
     x = self.per_layer_model_projection('...td,dnp->...tnp', x)
     x = self.per_layer_projection_norm(x)
+    if ignore_ple_tokens:
+      return x
     y = self.per_layer_input_embedding_table[(t,)]
     y *= jnp.sqrt(self.per_layer_input_dim).astype(y.dtype)
     return (x + y) * jax.lax.rsqrt(2.0).astype(x.dtype)

gemma/gm/nn/gemma4/_transformer.py

@@ -421,6 +421,7 @@ def _encode_and_get_inputs(
       attention_mask=None,
       positions=None,
       audio_soft_token_counts=None,
+      ignore_ple_tokens: bool = False,
   ) -> _Inputs:
     """Encode the text tokens, eventually including the vision embeddings."""
     if images is not None or audio is not None:
@@ -467,7 +468,9 @@ def _encode_and_get_inputs(
         )
 
       if self.config.per_layer_input_dim:
-        per_layer_inputs = self.embedder.encode_per_layer_input(x, tokens)
+        per_layer_inputs = self.embedder.encode_per_layer_input(
+            x, tokens, ignore_ple_tokens=ignore_ple_tokens
+        )
       else:
         per_layer_inputs = None
 
@@ -501,7 +504,7 @@ def _encode_and_get_inputs(
 
     if self.config.per_layer_input_dim:
       per_layer_inputs = self.embedder.encode_per_layer_input(
-          x, inputs.tokens_with_mm
+          x, inputs.tokens_with_mm, ignore_ple_tokens=ignore_ple_tokens
       )
     else:
       per_layer_inputs = None