Gemma3 initial commit

README.md

@@ -31,6 +31,7 @@ These are listed based on status and then alphabetically.
 | [ResNet50](bonsai/models/resnet50)            | Image classification | ✅     |         |
 | [VGG](bonsai/models/vgg19)                    | Image classification | ✅     |         |
 | [Dinov3](bonsai/models/dinov3)                | Vision FM            | ⚙️     |         |
+| [Gemma3](bonsai/models/gemma3)                | VLM                  | ⚙️     | Local attention cache and todos in file        |
 | [Mamba2](bonsai/models/mamba2)                | Language SSM         | ⚙️     | Caching and sharding        |
 | [umT5](bonsai/models/umt5)                    | LLM                  | ⚙️     | Caching and sharding        |
 | [ViT](bonsai/models/vit)                      | Image classification | ⚙️     | Sharding        |
@@ -39,7 +40,6 @@ These are listed based on status and then alphabetically.
 | [UNet](bonsai/models/unet/)                   | Image                | 🟡     | Need a reference implementation and numerical testing        |
 | [VAE](bonsai/models/vae/)                     | Generative model     | 🟡     | Need a reference implementation and numerical testing         |
 | [Whisper](bonsai/models/whisper/)             | Speech recognition   | 🟡     | Need more numerical testing and not all call methods implemented        |
-| Gemma3                                        |                      | ⏳     |         |
 | CLIP                                          |                      | ⏳     |         |
 
 

bonsai/models/gemma3/README.md

@@ -0,0 +1,24 @@
+# Gemma3 in JAX
+
+This directory contains a pure JAX implementation of the [Gemma3 model](https://deepmind.google/models/gemma/gemma-3/), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API. 
+
+Note that you need an access token to download the model weights. In order to run the scripts, make sure to save an environment variable `HF_TOKEN` with your huggingface access token. 
+
+
+## Model Configuration Support Status
+
+
+### Running this model
+
+
+```sh
+python3 -m bonsai.models.gemma3.tests.run_model
+```
+
+
+## How to contribute to this model
+
+### Remaining Tasks
+
+1. Update to include kv cache memory reduction benefits from local attention. Currently, decode generation is not performance optimized.
+2. Update to optimize parameter loading for larger models.