Add mamba2 model#103

Merged

jenriver merged 16 commits intojax-ml:mainfrom

CosmoNaught:add-mamba2-model

Dec 22, 2025

Commits on Dec 11, 2025

Add Mamba2 model implementation (JAX NNX)
CosmoNaught
committed

Commits on Dec 14, 2025

Merge branch 'main' into add-mamba2-model
CosmoNaught
authored

Commits on Dec 16, 2025

Commits on Dec 18, 2025

use broadcast cumsum in segsum to avoid OOM
CosmoNaught
committed

Commits on Dec 19, 2025

add golden parity tests against mamba_ssm reference
CosmoNaught
committed

Commits on Dec 22, 2025