Models quick import

bonsai/models/__init__.py

@@ -0,0 +1,59 @@
+from bonsai.models.convnext.modeling import ConvNeXt, ModelConfig as ConvNeXtConfig
+from bonsai.models.densenet121.modeling import DenseNet, ModelConfig as DenseNetConfig
+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.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
+from bonsai.models.resnet.modeling import ResNet, ModelConfig as ResNetConfig
+from bonsai.models.sam2.modeling import SAM2Base, SAM2ImagePredictor, ModelConfig as SAM2Config
+from bonsai.models.umt5.modeling import UMT5Model, ModelConfig as UMT5Config
+from bonsai.models.unet.modeling import UNet, ModelConfig as UNetConfig
+from bonsai.models.vae.modeling import VAE, ModelConfig as VAEConfig
+from bonsai.models.vgg19.modeling import VGG, ModelConfig as VGGConfig
+from bonsai.models.vit.modeling import ViTClassificationModel, ModelConfig as ViTClassificationModelConfig
+from bonsai.models.vjepa2.modeling import VJEPA2ForVideoClassification, VJEPA2Model, ModelConfig as VJEPA2Config
+from bonsai.models.whisper.modeling import Whisper, ModelConfig as WhisperConfig
+
+
+__all__ = [
+    "ConvNeXt",
+    "ConvNeXtConfig",
+    "DenseNet",
+    "DenseNetConfig",
+    "Dinov3ViTModel",
+    "Dinov3ViTModelConfig",
+    "EfficientNet",
+    "EfficientNetConfig",
+    "Gemma3Model",
+    "Gemma3ModelConfig",
+    "LLaDAModel",
+    "LLaDAModelConfig",
+    "Mamba2Config",
+    "Mamba2ForCausalLM",
+    "Mamba2Forecaster",
+    "Mamba2Model",
+    "Qwen3",
+    "Qwen3Config",
+    "ResNet",
+    "ResNetConfig",
+    "SAM2Base",
+    "SAM2Config",
+    "SAM2ImagePredictor",
+    "UMT5Config",
+    "UMT5Model",
+    "UNet",
+    "UNetConfig",
+    "VAE",
+    "VAEConfig",
+    "VGG",
+    "VGGConfig",
+    "ViTClassificationModel",
+    "ViTClassificationModelConfig",
+    "VJEPA2Config",
+    "VJEPA2ForVideoClassification",
+    "VJEPA2Model",
+    "Whisper",
+    "WhisperConfig",
+]

bonsai/models/mamba2/modeling.py

@@ -33,7 +33,7 @@
 
 
 @dataclasses.dataclass(frozen=True)
-class Mamba2Config:
+class ModelConfig:
     """Configuration for Mamba2 models."""
 
     vocab_size: int = 50280
@@ -96,14 +96,14 @@ def tree_unflatten(cls, aux_data, children):
 
 
 def create_empty_cache(
-    cfg: Mamba2Config,
+    cfg: ModelConfig,
     batch_size: int,
     dtype: jnp.dtype = jnp.float32,
 ) -> Mamba2Cache:
     """Create an empty cache for Mamba2 model.
 
     Args:
-        cfg: Mamba2Config for the model.
+        cfg: ModelConfig for the model.
         batch_size: Batch size for the cache.
         dtype: Data type for cache arrays.
 
@@ -304,7 +304,7 @@ def __call__(self, x: jnp.ndarray, conv_state: jnp.ndarray | None = None) -> tup
 class Mamba2Mixer(nnx.Module):
     """Mamba2 mixer block using the SSD algorithm."""
 
-    def __init__(self, cfg: Mamba2Config, layer_idx: int, *, rngs: nnx.Rngs):
+    def __init__(self, cfg: ModelConfig, layer_idx: int, *, rngs: nnx.Rngs):
         self.cfg = cfg
         self.layer_idx = layer_idx
         self.hidden_size = cfg.hidden_size
@@ -407,7 +407,7 @@ def __call__(
 class Mamba2Block(nnx.Module):
     """Single Mamba2 block with pre-norm and residual connection."""
 
-    def __init__(self, cfg: Mamba2Config, layer_idx: int, *, rngs: nnx.Rngs):
+    def __init__(self, cfg: ModelConfig, layer_idx: int, *, rngs: nnx.Rngs):
         self.cfg = cfg
         self.residual_in_fp32 = cfg.residual_in_fp32
         self.norm = RMSNorm(cfg.hidden_size, eps=cfg.layer_norm_epsilon, rngs=rngs)
@@ -430,7 +430,7 @@ def __call__(
 class Mamba2Model(nnx.Module):
     """Mamba2 backbone model (no task-specific head)."""
 
-    def __init__(self, cfg: Mamba2Config, *, rngs: nnx.Rngs):
+    def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs):
         self.cfg = cfg
         self.embedder = nnx.Embed(num_embeddings=cfg.vocab_size, features=cfg.hidden_size, rngs=rngs)
         self.layers = nnx.List([Mamba2Block(cfg, layer_idx=i, rngs=rngs) for i in range(cfg.num_hidden_layers)])
@@ -491,7 +491,7 @@ def __call__(
 class Mamba2ForCausalLM(nnx.Module):
     """Mamba2 model with causal language modeling head."""
 
-    def __init__(self, cfg: Mamba2Config, *, rngs: nnx.Rngs):
+    def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs):
         self.cfg = cfg
         self.backbone = Mamba2Model(cfg, rngs=rngs)
         if not cfg.tie_word_embeddings:
@@ -527,7 +527,7 @@ def from_pretrained(
         cls,
         model_id_or_path: str,
         *,
-        cfg: Mamba2Config | None = None,
+        cfg: ModelConfig | None = None,
         dtype: jnp.dtype = jnp.float32,
         seed: int = 0,
         revision: str = "main",
@@ -565,7 +565,7 @@ def __init__(
         self.output_dim = output_dim
 
         self.input_proj = nnx.Linear(input_dim, d_model, rngs=rngs)
-        cfg = Mamba2Config(
+        cfg = ModelConfig(
             vocab_size=1,
             hidden_size=d_model,
             state_size=d_state,

bonsai/models/mamba2/params.py

@@ -25,11 +25,11 @@
 from bonsai.models.mamba2 import modeling
 
 
-def create_random_model(cfg: modeling.Mamba2Config, seed: int = 0) -> modeling.Mamba2ForCausalLM:
+def create_random_model(cfg: modeling.ModelConfig, seed: int = 0) -> modeling.Mamba2ForCausalLM:
     """Create a randomly initialized Mamba2ForCausalLM.
 
     Args:
-        cfg: Mamba2Config for the model.
+        cfg: ModelConfig for the model.
         seed: Random seed for initialization.
 
     Returns:
@@ -263,7 +263,7 @@ def load_pytorch_weights(
 
 def create_model_from_torch_checkpoint(
     checkpoint_path: str,
-    cfg: modeling.Mamba2Config | None = None,
+    cfg: modeling.ModelConfig | None = None,
     dtype: jnp.dtype = jnp.float32,
     seed: int = 0,
 ) -> modeling.Mamba2ForCausalLM:
@@ -324,7 +324,7 @@ def create_model_from_torch_checkpoint(
 
 def create_model_from_huggingface(
     model_id: str,
-    cfg: modeling.Mamba2Config | None = None,
+    cfg: modeling.ModelConfig | None = None,
     dtype: jnp.dtype = jnp.float32,
     seed: int = 0,
     revision: str = "main",
@@ -342,7 +342,7 @@ def create_model_from_huggingface(
         Mamba2ForCausalLM with loaded weights.
 
     Example:
-        >>> cfg = modeling.Mamba2Config(
+        >>> cfg = modeling.ModelConfig(
         ...     vocab_size=50280, hidden_size=768,
         ...     state_size=128, num_hidden_layers=24, head_dim=64
         ... )
@@ -369,7 +369,7 @@ def create_model_from_huggingface(
         with open(config_path) as f:
             hf_config = json.load(f)
 
-        cfg = modeling.Mamba2Config(
+        cfg = modeling.ModelConfig(
             vocab_size=hf_config.get("vocab_size", 50280),
             hidden_size=hf_config.get("d_model", hf_config.get("hidden_size", 768)),
             state_size=hf_config.get("d_state", hf_config.get("state_size", 128)),

bonsai/models/mamba2/tests/run_model.py

@@ -101,7 +101,7 @@ def _greedy_generate_cached(
     prompt_ids_1d: jnp.ndarray,
     *,
     max_new_tokens: int,
-    cfg: modeling.Mamba2Config,
+    cfg: modeling.ModelConfig,
 ) -> jnp.ndarray:
     """Greedy generation with SSM state caching (O(n) complexity).
 
@@ -150,7 +150,7 @@ def run_model(*, max_new_tokens: int = 32, use_cache: bool = True) -> None:
         "What is the capital city of England?",
     ]
 
-    cfg = modeling.Mamba2Config(
+    cfg = modeling.ModelConfig(
         vocab_size=50288,
         hidden_size=768,
         state_size=128,

bonsai/models/mamba2/tests/test_outputs_mamba_2.py

@@ -28,30 +28,30 @@
 from bonsai.models.mamba2 import modeling, params
 
 
-class TestMamba2Config(absltest.TestCase):
-    """Tests for Mamba2Config."""
+class TestModelConfig(absltest.TestCase):
+    """Tests for ModelConfig."""
 
     def test_default_config(self):
         """Test default config values."""
-        cfg = modeling.Mamba2Config()
+        cfg = modeling.ModelConfig()
         self.assertEqual(cfg.vocab_size, 50280)
         self.assertEqual(cfg.hidden_size, 768)
         self.assertEqual(cfg.num_hidden_layers, 24)
 
     def test_intermediate_size(self):
         """Test intermediate_size property."""
-        cfg = modeling.Mamba2Config(hidden_size=512, expand=2)
+        cfg = modeling.ModelConfig(hidden_size=512, expand=2)
         self.assertEqual(cfg.intermediate_size, 1024)
 
     def test_num_heads(self):
         """Test num_heads property."""
-        cfg = modeling.Mamba2Config(hidden_size=512, expand=2, head_dim=64)
+        cfg = modeling.ModelConfig(hidden_size=512, expand=2, head_dim=64)
         # intermediate_size = 1024, head_dim = 64 -> num_heads = 16
         self.assertEqual(cfg.num_heads, 16)
 
     def test_predefined_configs(self):
         """Test predefined configuration methods."""
-        cfg_tiny = modeling.Mamba2Config.tiny()
+        cfg_tiny = modeling.ModelConfig.tiny()
         self.assertEqual(cfg_tiny.hidden_size, 64)
         self.assertEqual(cfg_tiny.num_hidden_layers, 2)
 
@@ -155,7 +155,7 @@ class TestMamba2Model(absltest.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.cfg = modeling.Mamba2Config.tiny()
+        self.cfg = modeling.ModelConfig.tiny()
         self.model = modeling.Mamba2Model(self.cfg, rngs=nnx.Rngs(42))
 
     def test_output_shape(self):
@@ -206,7 +206,7 @@ class TestMamba2ForCausalLM(absltest.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.cfg = modeling.Mamba2Config.tiny()
+        self.cfg = modeling.ModelConfig.tiny()
         self.model = modeling.Mamba2ForCausalLM(self.cfg, rngs=nnx.Rngs(42))
 
     def test_output_shape(self):
@@ -267,7 +267,7 @@ class TestParameters(absltest.TestCase):
 
     def test_create_random_model(self):
         """Test random model creation."""
-        cfg = modeling.Mamba2Config.tiny()
+        cfg = modeling.ModelConfig.tiny()
         model = params.create_random_model(cfg, seed=42)
         self.assertIsInstance(model, modeling.Mamba2ForCausalLM)
 
@@ -292,7 +292,7 @@ class TestJIT(absltest.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.cfg = modeling.Mamba2Config.tiny()
+        self.cfg = modeling.ModelConfig.tiny()
 
     def test_jit_backbone(self):
         """Test that backbone can be JIT compiled."""
@@ -321,7 +321,7 @@ class TestGradients(absltest.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.cfg = modeling.Mamba2Config.tiny()
+        self.cfg = modeling.ModelConfig.tiny()
 
     def test_gradients_exist(self):
         """Test that gradients can be computed."""
@@ -364,7 +364,7 @@ def setUpClass(cls):
 
     def test_hidden_state_parity(self):
         """Test last_hidden_state matches mamba_ssm reference within numerical tolerance."""
-        cfg = modeling.Mamba2Config(
+        cfg = modeling.ModelConfig(
             vocab_size=50288,
             hidden_size=768,
             state_size=128,
@@ -388,7 +388,7 @@ def test_hidden_state_parity(self):
 
     def test_logits_parity(self):
         """Test logits match mamba_ssm reference within numerical tolerance."""
-        cfg = modeling.Mamba2Config(
+        cfg = modeling.ModelConfig(
             vocab_size=50288,
             hidden_size=768,
             state_size=128,
@@ -416,7 +416,7 @@ class TestMamba2Cache(absltest.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.cfg = modeling.Mamba2Config.tiny()
+        self.cfg = modeling.ModelConfig.tiny()
         self.model = modeling.Mamba2ForCausalLM(self.cfg, rngs=nnx.Rngs(42))
 
     def test_cache_shapes(self):

bonsai/models/sam2/modeling.py

@@ -135,7 +135,7 @@ class MemoryEncoderConfig:
 
 
 @dataclass(frozen=True)
-class SAM2Config:
+class ModelConfig:
     image_encoder: ImageEncoderConfig
     memory_attention: MemoryAttentionConfig
     memory_encoder: MemoryEncoderConfig
@@ -1229,7 +1229,7 @@ def _apply_non_overlapping_constraints(
         return jnp.where(keep, pred_masks, clamped)
 
 
-def build_sam2_model_from_config(cfg: SAM2Config, rngs: nnx.Rngs) -> SAM2Base:
+def build_sam2_model_from_config(cfg: ModelConfig, rngs: nnx.Rngs) -> SAM2Base:
     # === Position Encodings ===
     pos_enc_backbone = PositionEmbeddingSine(
         num_pos_feats=cfg.image_encoder.neck.position_encoding.num_pos_feats,

bonsai/models/sam2/params.py

@@ -699,7 +699,7 @@ def parse_key(k):
 
 def create_sam2_from_pretrained(
     file_path: str,
-    config: model_lib.SAM2Config,
+    config: model_lib.ModelConfig,
     *,
     mesh: jax.sharding.Mesh | None = None,
 ):

bonsai/models/sam2/tests/SAM2_image_predictor_example.ipynb

@@ -327,7 +327,7 @@
     "MODEL_CP_PATH = \"./checkpoints/\" + model_name.split(\"/\")[1]\n",
     "snapshot_download(model_name, local_dir=MODEL_CP_PATH)\n",
     "\n",
-    "config = modeling.SAM2Config.sam2_small()\n",
+    "config = modeling.ModelConfig.sam2_small()\n",
     "model_obj = params.create_sam2_from_pretrained(MODEL_CP_PATH + \"/model.safetensors\", config)\n",
     "\n",
     "predictor = modeling.SAM2ImagePredictor(model_obj)"

bonsai/models/sam2/tests/SAM2_image_predictor_example.md

@@ -138,7 +138,7 @@ model_name = "facebook/sam2-hiera-small-hf"
 MODEL_CP_PATH = "./checkpoints/" + model_name.split("/")[1]
 snapshot_download(model_name, local_dir=MODEL_CP_PATH)
 
-config = modeling.SAM2Config.sam2_small()
+config = modeling.ModelConfig.sam2_small()
 model_obj = params.create_sam2_from_pretrained(MODEL_CP_PATH + "/model.safetensors", config)
 
 predictor = modeling.SAM2ImagePredictor(model_obj)

bonsai/models/sam2/tests/run_model.py

@@ -63,7 +63,7 @@ def run_model(MODEL_CP_PATH=None):
         convert_pt_to_safetensors(pt_path, safetensors_path)
 
     # 2. Create SAM2 model and predictor
-    config = modeling.SAM2Config.sam2_tiny()
+    config = modeling.ModelConfig.sam2_tiny()
     model_obj = params.create_sam2_from_pretrained(safetensors_path, config)
     model = modeling.SAM2ImagePredictor(model_obj)
 

bonsai/models/umt5/README.md

@@ -26,5 +26,5 @@ python3 -m bonsai.models.umt5.tests.run_model
 ## How to contribute to this model
 
 We welcome contributions! You can contribute to this model via the following:
-* Add a model config variant from the above `🟡 Not started` to `class UMT5Config` in [modeling.py](modeling.py). Make sure your code is runnable on at least one hardware before creating a PR.
+* Add a model config variant from the above `🟡 Not started` to `class ModelConfig` in [modeling.py](modeling.py). Make sure your code is runnable on at least one hardware before creating a PR.
 * Got some hardware? Run [run_model.py](tests/run_model.py) the existing configs above on hardwares marked `❔ Needs check`. Mark as `✅ Runs` or `⛔️ Not supported`.