gm.nn.Gemma4_E2B

class gemma.gm.nn.Gemma4_E2B(

config: _config.TransformerConfig = TransformerConfig(num_embed=262144,

embed_dim=1536,

hidden_dim=6144,

num_heads=8,

head_dim=256,

num_kv_heads=1,

final_logit_softcap=30.0,

use_post_attn_norm=True,

use_post_ffw_norm=True,

attention_types=(<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.GLOBAL: 1>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.GLOBAL: 1>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.GLOBAL: 1>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.GLOBAL: 1>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.GLOBAL: 1>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.GLOBAL: 1>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.LOCAL_SLIDING: 2>,

<AttentionType.GLOBAL: 1>),

attn_logits_soft_cap=None,

sliding_window_size=512,

qk_norm_with_scale=True,

num_global_kv_heads=None,

global_key_size=512,

k_eq_v_global=False,

global_rope_proportion=0.25,

local_rope_proportion=1.0,

local_base_frequency=10000,

global_base_frequency=1000000,

local_scale_factor=1.0,

global_scale_factor=1.0,

per_layer_input_dim=256,

kv_cache_sharing_config=KVCacheSharingConfig(frac_shared_layers=0.5714285714285714,

share_global=True,

share_local=True),

override_kv_shared_ffw_hidden=12288,

vision_encoder=VisionEncoder(     # attributes     d_model = 768     num_layers = 16     num_heads = 12     ffw_hidden = 3072     patch_size = 16     output_length = 280     pos_emb_shape_yx = (10240,

2)     pooling_kernel_size=3     use_clipped_linears = True     standardize_embeddings = False ),

audio_encoder=ConformerConfig(num_layers=12,

model_dims=1024,

lm_model_dims=1536,

atten_num_heads=8,

atten_left_context=13,

atten_right_context=0,

conv_kernel_size=5,

gradient_clipping=10000000000.0,

conf_reduction_factor=1,

param_dtype=<class 'jax.numpy.float32'>,

compute_dtype=None),

enable_moe=False,

num_experts=0,

expert_dim=0,

top_k_experts=8,

moe_dense_hidden_dim=0,

use_bidirectional_attention=None),

text_only: bool = True,

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

name: str | None = None,

*,

return_last_only: bool | None = None,

dtype: jnp.dtype = <class 'jax.numpy.bfloat16'>,

tokens: kontext.Key = '__KEY_REQUIRED__',

images: kontext.Key | None = None,

)

Bases: gemma.gm.nn.gemma4._gemma4._Gemma4Base

Gemma 4 E2B model.

attention_pattern = (AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL)

global_local_pattern = (AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL)

config: gemma.gm.nn.gemma4._config.TransformerConfig = TransformerConfig(num_embed=262144, embed_dim=1536, hidden_dim=6144, num_heads=8, head_dim=256, num_kv_heads=1, final_logit_softcap=30.0, use_post_attn_norm=True, use_post_ffw_norm=True, attention_types=(<AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.GLOBAL: 1>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.GLOBAL: 1>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.GLOBAL: 1>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.GLOBAL: 1>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.GLOBAL: 1>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.GLOBAL: 1>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.LOCAL_SLIDING: 2>, <AttentionType.GLOBAL: 1>), attn_logits_soft_cap=None, sliding_window_size=512, qk_norm_with_scale=True, num_global_kv_heads=None, global_key_size=512, k_eq_v_global=False, global_rope_proportion=0.25, local_rope_proportion=1.0, local_base_frequency=10000, global_base_frequency=1000000, local_scale_factor=1.0, global_scale_factor=1.0, per_layer_input_dim=256, kv_cache_sharing_config=KVCacheSharingConfig(frac_shared_layers=0.5714285714285714, share_global=True, share_local=True), override_kv_shared_ffw_hidden=12288, vision_encoder=VisionEncoder(     # attributes     d_model = 768     num_layers = 16     num_heads = 12     ffw_hidden = 3072     patch_size = 16     output_length = 280     pos_emb_shape_yx = (10240, 2)     pooling_kernel_size = 3     use_clipped_linears = True     standardize_embeddings = False ), audio_encoder=ConformerConfig(num_layers=12, model_dims=1024, lm_model_dims=1536, atten_num_heads=8, atten_left_context=13, atten_right_context=0, conv_kernel_size=5, gradient_clipping=10000000000.0, conf_reduction_factor=1, param_dtype=<class 'jax.numpy.float32'>, compute_dtype=None), enable_moe=False, num_experts=0, expert_dim=0, top_k_experts=8, moe_dense_hidden_dim=0, use_bidirectional_attention=None)

INFO: ClassVar[gemma.gm.nn.gemma4._transformer.ModelInfo] = ModelInfo(tokenizer_version=4, default_ckpt=None)

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