Add Llama3.2 (1B & 3B) - #138

Open

Moriyuki-S wants to merge 22 commits into

jax-ml:mainfrom

Moriyuki-S:feat/implement-llama

Commits on Jan 22, 2026

Commits on Jan 23, 2026

Commits on Jan 30, 2026

docs(llama3.2): update Flax NNX API link to stable version
Moriyuki-S
committed
c5603bd
Merge branch 'main' into feat/implement-llama
jenriver
authored
1dc2afc

Commits on Feb 1, 2026

fix(llama3.2): correct sharding axes for GQA and fix explicit mesh errors
Moriyuki-S
committed
b4cf559
docs(llama3.2): enhance docstring for compute_positions_from_segment_ids to clarify packed sequences support
Moriyuki-S
committed
a678b01
test(llama3.2): consolidate test configs
Moriyuki-S
committed
ee9e595
test(llama3.2): refactor padding tests
Moriyuki-S
committed
4c8ba67
test(llama3.2): add attention mask tests for padding and future tokens
Moriyuki-S
committed
7407439
refactor(llama3.2): enhance run_model script with argument parsing and remove deprecated base model test
Moriyuki-S
committed
fee4a82
Merge branch 'feat/implement-llama' of https://github.com/Moriyuki-S/bonsai into feat/implement-llama
Moriyuki-S
committed
f2d3a8d

Commits on Feb 11, 2026

refactor(llama3.2): add no mesh support
Moriyuki-S
committed
c047b0b
style(llama3.2): clean up code formatting
Moriyuki-S
committed
a32f929
style(llama3.2): rename test classes
Moriyuki-S
committed
c8404a0
Merge branch 'main' into feat/implement-llama

jenriver
authored and
Moriyuki-S
committed
d6a9d40
Merge branch 'main' into feat/implement-llama
Moriyuki-S
committed
6c13f56