gm.nn.Embedder

class gemma.gm.nn.Embedder(

vocab_size: int,

embed_dim: int,

vision_proj_dim: int | None = None,

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>,

name: str | None = None,

)

Bases: flax.linen.module.Module

Embedder module.

vocab_size: int

embed_dim: int

vision_proj_dim: int | None = None

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

encode(x: jax.Array) → jax.Array

Encodes the input tokens.

Parameters:

x – Input tokens of shape [seq_len] or [batch_size, seq_len], where each token is an integer in [0, vocab_size).

Returns:

Encoded tokens of shape [seq_len, embed_dim] or [batch_size, seq_len, embed_dim].

decode(x: jax.Array) → jax.Array

Decodes the input vectors.

Parameters:

x – Array of shape [seq_len, embed_dim] or [batch_size, seq_len, embed_dim].

Returns:

Array of shape [seq_len, vocab_size] or [batch_size, seq_len, vocab_size].

encode_vision(x: jax.Array) → jax.Array

Projects siglip embeddings to the embedding space of the text encoder.

name: str | None = None

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None

scope: flax.core.scope.Scope | None = None