Gemma 4

README.md

@@ -32,6 +32,7 @@ These are listed based on status and then alphabetically.
 | [VGG](bonsai/models/vgg19)                    | Image classification | ✅     |         |
 | [Dinov3](bonsai/models/dinov3)                | Vision FM            | ⚙️     |         |
 | [Gemma3](bonsai/models/gemma3)                | VLM                  | ⚙️     | Local attention cache and todos in file        |
+| [Gemma4](bonsai/models/gemma4)                | LLM / MoE            | ✅     |                                                |
 | [Mamba2](bonsai/models/mamba2)                | Language SSM         | ⚙️     | Caching and sharding        |
 | [umT5](bonsai/models/umt5)                    | LLM                  | ⚙️     | Caching and sharding        |
 | [ViT](bonsai/models/vit)                      | Image classification | ⚙️     | Sharding        |

bonsai/models/__init__.py

@@ -3,6 +3,7 @@
 from bonsai.models.dinov3.modeling import Dinov3ViTModel, ModelConfig as Dinov3ViTModelConfig
 from bonsai.models.efficientnet.modeling import EfficientNet, ModelConfig as EfficientNetConfig
 from bonsai.models.gemma3.modeling import Gemma3Model, ModelConfig as Gemma3ModelConfig
+from bonsai.models.gemma4.modeling import Gemma4Model, Gemma4ForCausalLM, ModelConfig as Gemma4Config
 from bonsai.models.llada.modeling import LLaDAModel, ModelConfig as LLaDAModelConfig
 from bonsai.models.mamba2.modeling import Mamba2ForCausalLM, Mamba2Forecaster, Mamba2Model, ModelConfig as Mamba2Config
 from bonsai.models.qwen3.modeling import Qwen3, ModelConfig as Qwen3Config
@@ -29,6 +30,9 @@
     "EfficientNetConfig",
     "Gemma3Model",
     "Gemma3ModelConfig",
+    "Gemma4Config",
+    "Gemma4ForCausalLM",
+    "Gemma4Model",
     "LLaDAModel",
     "LLaDAModelConfig",
     "Mamba2Config",

bonsai/models/dinov3/modeling.py

@@ -351,6 +351,6 @@ def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
         return params.create_model_from_safe_tensors(model_ckpt_path, config)
 
 
-@jax.jit()
+@jax.jit
 def forward(model: Dinov3ViTModel, inputs: Array):
     return model(inputs)

bonsai/models/gemma4/README.md

@@ -0,0 +1,41 @@
+# Gemma 4 Model
+
+This is the implementation of the Gemma 4 architecture in JAX and Flax NNX for `bonsai`.
+
+## Architecture Details
+
+Gemma 4 introduces a hybrid MoE (Mixture of Experts) and hybrid attention pattern:
+- **Attention Pattern**: A repeating cycle of 5 Local Sliding Window attention layers followed by 1 Global attention layer (`LOCAL_SLIDING`, `LOCAL_SLIDING`, `LOCAL_SLIDING`, `LOCAL_SLIDING`, `LOCAL_SLIDING`, `GLOBAL`).
+- **Mixture of Experts**: Combines a top-k routed expert module with a single persistent, wider shared expert module.
+- **Normalization**: Utilizes specialized zero-scale RMSNorm layers within MoE gating mechanisms alongside standard offset-scale RMSNorms (`1 + scale`) throughout the model. Furthermore, Query and Key embeddings use RMSNorm.
+- **Logit Soft-capping**: Output logits are optionally soft-capped (usually value 30.0) before softmax.
+
+## Configuration
+
+The default base configuration configures:
+- Hybrid attention logic with independent relative RoPE frequency parameters (`global_rope_proportion`, `local_rope_proportion`).
+- The necessary layer size parameters: `num_hidden_layers`, `hidden_size`, `intermediate_size`.
+- The MoE parameters: `num_experts`, `num_shared_experts`, and `num_experts_per_tok`.
+
+## Example Usage
+
+```python
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from bonsai.models.gemma4 import Gemma4Config, Gemma4ForCausalLM
+
+# Initialize base configuration
+config = Gemma4Config.gemma4_base()
+
+# Initialize model
+rngs = nnx.Rngs(0)
+model = Gemma4ForCausalLM(config, rngs=rngs)
+
+# Forward pass
+input_ids = jnp.array([[1, 2, 3, 4]])
+positions = jnp.array([[0, 1, 2, 3]])
+
+logits = model(input_ids, positions=positions)
+print(logits.shape) # (1, 4, 256000)
+```

bonsai/models/gemma4/__init__.py

@@ -0,0 +1,35 @@
+# Copyright 2026 The JAX Authors.
+#
+# 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
+#
+#     http://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.
+
+from bonsai.models.gemma4.modeling import (
+    Gemma4Model,
+    Gemma4ForCausalLM,
+    ModelConfig as Gemma4Config,
+    LayerCache,
+    Cache,
+    init_cache,
+    forward,
+)
+from bonsai.models.gemma4.params import create_gemma4_from_pretrained
+
+__all__ = [
+    "Gemma4Model",
+    "Gemma4ForCausalLM",
+    "Gemma4Config",
+    "LayerCache",
+    "Cache",
+    "init_cache",
+    "forward",
+    "create_gemma4_from_pretrained",
+]