Skip to content

AIMv2

AIMv2 is a vision transformer trained autoregressively on image-text pairs. Rather than a contrastive or reconstruction objective, it predicts image tokens in a causal left-to-right sequence conditioned on a multimodal prefix, producing strong general-purpose patch-level representations.

Paper: "Multimodal Autoregressive Pre-training of Large Vision Encoders" (El-Nouby et al., CVPR 2025) Code: github.com/apple/ml-aim

The jimm implementation returns all patch token embeddings after the final RMSNorm with shape (batch, n_patches, hidden_size). AIMv2 has no CLS token; downstream tasks pool over patches or select specific tokens.

Key architectural differences from standard ViT:

Feature ViT AIMv2
Normalization LayerNorm RMSNorm
MLP Standard GELU SwiGLU
Pre-norm placement After patch + pos On patches only, before pos
CLS token ✅ ❌
Output CLS embedding (B, D) All patches (B, N, D)
Attention / MLP biases Yes No¹

¹ Attention and MLP projections have no bias. The patch embedding convolution does have a bias (use_patch_bias=True), consistent with the HuggingFace checkpoint.

Supported models

All AIMv2 models use patch size 14 and 24 transformer layers regardless of scale. The four architecture sizes scale width only.

Large - 307M params

HuggingFace ID hidden_size Heads intermediate_size Image size
apple/aimv2-large-patch14-224 1024 8 2816 224
apple/aimv2-large-patch14-336 1024 8 2816 336
apple/aimv2-large-patch14-448 1024 8 2816 448
apple/aimv2-large-patch14-224-distilled 1024 8 2816 224
apple/aimv2-large-patch14-224-lit 1024 8 2816 224

Huge - 680M params

HuggingFace ID hidden_size Heads intermediate_size Image size
apple/aimv2-huge-patch14-224 1536 12 4096 224
apple/aimv2-huge-patch14-336 1536 12 4096 336
apple/aimv2-huge-patch14-448 1536 12 4096 448

1B - 1.2B params

HuggingFace ID hidden_size Heads intermediate_size Image size
apple/aimv2-1B-patch14-224 2048 16 5632 224
apple/aimv2-1B-patch14-336 2048 16 5632 336
apple/aimv2-1B-patch14-448 2048 16 5632 448

3B - 2.7B params

HuggingFace ID hidden_size Heads intermediate_size Image size
apple/aimv2-3B-patch14-224 3072 24 8192 224
apple/aimv2-3B-patch14-336 3072 24 8192 336
apple/aimv2-3B-patch14-448 3072 24 8192 448

Note: apple/aimv2-large-patch14-native uses sinusoidal position embeddings for variable resolution and is not currently supported. The lit model is a full multimodal checkpoint (vision + text encoder + projectors); AIMv2Model.from_pretrained automatically extracts just the vision encoder from it.

Basic usage

import jimm
import numpy as np

model = jimm.AIMv2Model.from_pretrained("apple/aimv2-large-patch14-224")

images = np.random.rand(4, 224, 224, 3).astype(np.float32)
patch_embeddings = model(images)  # shape: (4, 256, 1024)

The model returns all 256 patch tokens (16×16 grid) for a 224×224 image with patch size 14. To get a single image embedding, pool over patches:

# Mean pool across patch dimension
image_embeddings = patch_embeddings.mean(axis=1)  # shape: (4, 1024)

Flash / Splash Attention

AIMv2 supports hardware-accelerated attention via Tokamax:

Backend Hardware Notes
"mosaic" NVIDIA H100 (SM90) / B100 (SM100) Pallas Mosaic GPU kernel
"triton" Any NVIDIA GPU Pallas Triton kernel
"cudnn" NVIDIA GPU Via JAX-NN / cuDNN
"mosaic_tpu" TPU v5+ (all generations) Splash attention (block-sparse)
"xla_chunked" GPU / TPU Flash-style chunked XLA
"xla" Any Standard XLA fallback 
import jimm

# GPU: try H100 Mosaic kernel, fall back to Triton, then XLA
model = jimm.AIMv2Model.from_pretrained("apple/aimv2-large-patch14-448",
                                         attention_fn=jimm.make_tokamax_attention(["mosaic", "triton", "xla"]))

# TPU: try Splash attention, fall back to chunked XLA
model = jimm.AIMv2Model.from_pretrained("apple/aimv2-large-patch14-448",
                                         attention_fn=jimm.make_tokamax_attention(["mosaic_tpu", "xla_chunked"]))

Note: AIMv2-Large produces 256 patch tokens at 224×224, 576 at 336×336, and 1024 at 448×448. Flash attention is most beneficial at the larger image sizes.

FSDP / Explicit Sharding

AIMv2 supports JAX explicit sharding (FSDP-style) via AIMv2Sharding, which inherits ViTSharding with no overrides needed (AIMv2 has no LayerScale, no register tokens, and no MLP bias tensors to shard).

from jax.experimental import mesh_utils
from jax.sharding import AxisType, Mesh
import jax

n_devices = jax.device_count()
mesh = Mesh(
    mesh_utils.create_device_mesh((1, n_devices)),
    ("data", "fsdp"),
    axis_types=(AxisType.Explicit, AxisType.Explicit),
)
jax.set_mesh(mesh)

model = jimm.AIMv2Model.from_pretrained("apple/aimv2-large-patch14-224")
# model params are automatically sharded across fsdp axis

To disable sharding, pass sharding=NoSharding() (import: from jimm.common.sharding import NoSharding).

jimm.models.aimv2.AIMv2Model

Bases: Module

AIMv2 vision encoder.

Implements the AIMv2 vision transformer with RMSNorm, SwiGLU MLP, no attention/MLP biases, and absolute learnable positional embeddings. Returns all patch token embeddings (no CLS token).

Source code in src/jimm/models/aimv2/aimv2_model.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
class AIMv2Model(nnx.Module):
    """AIMv2 vision encoder.

    Implements the AIMv2 vision transformer with RMSNorm, SwiGLU MLP, no
    attention/MLP biases, and absolute learnable positional embeddings.
    Returns all patch token embeddings (no CLS token).
    """

    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 14,
        in_channels: int = 3,
        hidden_size: int = 1024,
        num_layers: int = 24,
        num_heads: int = 8,
        mlp_dim: int = 2816,
        rms_norm_eps: float = 1e-5,
        use_gradient_checkpointing: bool = False,
        attention_fn: Callable[..., Any] | None = None,
        rngs: rnglib.Rngs | None = None,
        dtype: DTypeLike = jnp.float32,
        param_dtype: DTypeLike = jnp.float32,
        sharding: ShardingSpec | None = None,
    ) -> None:
        """Initialize AIMv2Model.

        Args:
            img_size (int, optional): Input image size. Defaults to 224.
            patch_size (int, optional): Patch size. Defaults to 14.
            in_channels (int, optional): Number of input channels. Defaults to 3.
            hidden_size (int, optional): Hidden dimension. Defaults to 1024 (Large).
            num_layers (int, optional): Number of transformer layers. Defaults to 24 (Large).
            num_heads (int, optional): Number of attention heads. Defaults to 8 (Large).
            mlp_dim (int, optional): MLP intermediate dimension. Defaults to 2816 (Large).
            rms_norm_eps (float, optional): RMSNorm epsilon. Defaults to 1e-5.
            use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
            attention_fn (Callable[..., Any] | None, optional): Custom attention function. Defaults to None.
            rngs (rnglib.Rngs | None, optional): RNG state. If None, initializes to nnx.Rngs(0).
            dtype (DTypeLike, optional): Computation dtype. Defaults to jnp.float32.
            param_dtype (DTypeLike, optional): Parameter dtype. Defaults to jnp.float32.
            sharding (ShardingSpec | None, optional): Sharding specification. Defaults to AIMv2Sharding().
        """
        if rngs is None:
            rngs = nnx.Rngs(0)
        if sharding is None:
            sharding = AIMv2Sharding()
        self._original_config = None
        self.encoder = VisionTransformerBase(
            img_size=img_size,
            patch_size=patch_size,
            in_channels=in_channels,
            hidden_size=hidden_size,
            num_layers=num_layers,
            num_heads=num_heads,
            mlp_dim=mlp_dim,
            layernorm_epsilon=rms_norm_eps,
            pooling_type="ALL",
            dropout_rate=0.0,
            act_fn=jax.nn.silu,
            use_pre_norm=True,
            pre_norm_before_pos=True,
            use_patch_bias=True,
            use_gradient_checkpointing=use_gradient_checkpointing,
            use_layer_scale=False,
            use_rope=False,
            num_register_tokens=0,
            use_gated_mlp=True,
            key_bias=False,
            attn_bias=False,
            mlp_bias=False,
            use_rms_norm=True,
            attention_fn=attention_fn,
            rngs=rngs,
            dtype=dtype,
            param_dtype=param_dtype,
            sharding=sharding,
        )

    def __call__(self, x: Float[Array, "batch height width channels"]) -> Float[Array, "batch n_patches hidden_size"]:
        """Run AIMv2 forward pass and return all patch token embeddings.

        Args:
            x (Float[Array, "batch height width channels"]): Batch of input images in BHWC format.
                Height and width must each be divisible by patch_size.

        Returns:
            Float[Array, "batch n_patches hidden_size"]: Patch token embeddings after final RMSNorm.
        """
        return self.encoder(x)

    def save_pretrained(self, save_directory: str) -> None:
        """Save model weights and config in HuggingFace format.

        Args:
            save_directory (str): Directory path where the model will be saved.
        """
        from .params import save_pretrained as _save_pretrained

        _save_pretrained(self, save_directory)

    @classmethod
    def from_pretrained(
        cls,
        model_name_or_path: str,
        use_pytorch: bool = False,
        rngs: rnglib.Rngs | None = None,
        dtype: DTypeLike = jnp.float32,
        param_dtype: DTypeLike = jnp.float32,
        sharding: ShardingSpec | None = None,
        use_gradient_checkpointing: bool = False,
        attention_fn: Callable[..., Any] | None = None,
    ) -> "AIMv2Model":
        """Load a pretrained AIMv2 model from a local path or HuggingFace Hub.

        Args:
            model_name_or_path (str): Local directory or HuggingFace model ID
                (e.g. "apple/aimv2-large-patch14-224").
            use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
            rngs (rnglib.Rngs | None): RNG state. If None, initializes to nnx.Rngs(0).
            dtype (DTypeLike): Computation dtype. Defaults to jnp.float32.
            param_dtype (DTypeLike): Parameter dtype. Defaults to jnp.float32.
            sharding (ShardingSpec | None): Sharding specification. Defaults to AIMv2Sharding().
            use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
            attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

        Returns:
            AIMv2Model: Model with pretrained weights loaded.
        """
        from .params import load_from_pretrained

        return load_from_pretrained(
            cls,
            model_name_or_path,
            use_pytorch,
            rngs,
            dtype,
            param_dtype,
            sharding,
            use_gradient_checkpointing,
            attention_fn,
        )

    @classmethod
    def _parse_config(cls, config: dict[str, Any]) -> dict[str, Any]:
        if config.get("model_type") == "aimv2":
            config = config["vision_config"]
        return {
            "img_size": config.get("image_size", 224),
            "patch_size": config.get("patch_size", 14),
            "in_channels": config.get("num_channels", 3),
            "hidden_size": config["hidden_size"],
            "num_layers": config["num_hidden_layers"],
            "num_heads": config["num_attention_heads"],
            "mlp_dim": config["intermediate_size"],
            "rms_norm_eps": config.get("rms_norm_eps", 1e-5),
        }

    @classmethod
    def from_config(
        cls,
        config: dict[str, Any],
        *,
        rngs: rnglib.Rngs | None = None,
        dtype: DTypeLike = jnp.float32,
        param_dtype: DTypeLike = jnp.float32,
        sharding: ShardingSpec | None = None,
        use_gradient_checkpointing: bool = False,
        attention_fn: Callable[..., Any] | None = None,
    ) -> "AIMv2Model":
        """Create AIMv2Model from a HuggingFace-compatible config dict.

        Args:
            config (dict[str, Any]): Configuration dictionary in HuggingFace AIMv2 format.
            rngs (rnglib.Rngs | None): RNG state. If None, initializes to nnx.Rngs(0).
            dtype (DTypeLike): Computation dtype. Defaults to jnp.float32.
            param_dtype (DTypeLike): Parameter dtype. Defaults to jnp.float32.
            sharding (ShardingSpec | None): Sharding specification. Defaults to AIMv2Sharding().
            use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
            attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

        Returns:
            AIMv2Model: Model with randomly initialized weights matching the given config.
        """
        if rngs is None:
            rngs = nnx.Rngs(0)
        return cls(
            **cls._parse_config(config),
            use_gradient_checkpointing=use_gradient_checkpointing,
            attention_fn=attention_fn,
            rngs=rngs,
            dtype=dtype,
            param_dtype=param_dtype,
            sharding=sharding,
        )

__call__(x)

Run AIMv2 forward pass and return all patch token embeddings.

Parameters:

Name Type Description Default
x Float[Array, 'batch height width channels']

Batch of input images in BHWC format. Height and width must each be divisible by patch_size.

 required

Returns:

Type Description
Float[Array, 'batch n_patches hidden_size']

Float[Array, "batch n_patches hidden_size"]: Patch token embeddings after final RMSNorm.
Source code in src/jimm/models/aimv2/aimv2_model.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
def __call__(self, x: Float[Array, "batch height width channels"]) -> Float[Array, "batch n_patches hidden_size"]:
    """Run AIMv2 forward pass and return all patch token embeddings.

    Args:
        x (Float[Array, "batch height width channels"]): Batch of input images in BHWC format.
            Height and width must each be divisible by patch_size.

    Returns:
        Float[Array, "batch n_patches hidden_size"]: Patch token embeddings after final RMSNorm.
    """
    return self.encoder(x)

__init__(img_size=224, patch_size=14, in_channels=3, hidden_size=1024, num_layers=24, num_heads=8, mlp_dim=2816, rms_norm_eps=1e-05, use_gradient_checkpointing=False, attention_fn=None, rngs=None, dtype=jnp.float32, param_dtype=jnp.float32, sharding=None)

Initialize AIMv2Model.

Parameters:

Name Type Description Default
img_size int

Input image size. Defaults to 224.

 224
patch_size int

Patch size. Defaults to 14.

 14
in_channels int

Number of input channels. Defaults to 3.

 3
hidden_size int

Hidden dimension. Defaults to 1024 (Large).

 1024
num_layers int

Number of transformer layers. Defaults to 24 (Large).

 24
num_heads int

Number of attention heads. Defaults to 8 (Large).

 8
mlp_dim int

MLP intermediate dimension. Defaults to 2816 (Large).

 2816
rms_norm_eps float

RMSNorm epsilon. Defaults to 1e-5.

 1e-05
use_gradient_checkpointing bool

Whether to use gradient checkpointing. Defaults to False.

 False
attention_fn Callable[..., Any] | None

Custom attention function. Defaults to None.

 None
rngs Rngs | None

RNG state. If None, initializes to nnx.Rngs(0).

 None
dtype DTypeLike

Computation dtype. Defaults to jnp.float32.

 float32
param_dtype DTypeLike

Parameter dtype. Defaults to jnp.float32.

 float32
sharding ShardingSpec | None

Sharding specification. Defaults to AIMv2Sharding().

 None
Source code in src/jimm/models/aimv2/aimv2_model.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def __init__(
    self,
    img_size: int = 224,
    patch_size: int = 14,
    in_channels: int = 3,
    hidden_size: int = 1024,
    num_layers: int = 24,
    num_heads: int = 8,
    mlp_dim: int = 2816,
    rms_norm_eps: float = 1e-5,
    use_gradient_checkpointing: bool = False,
    attention_fn: Callable[..., Any] | None = None,
    rngs: rnglib.Rngs | None = None,
    dtype: DTypeLike = jnp.float32,
    param_dtype: DTypeLike = jnp.float32,
    sharding: ShardingSpec | None = None,
) -> None:
    """Initialize AIMv2Model.

    Args:
        img_size (int, optional): Input image size. Defaults to 224.
        patch_size (int, optional): Patch size. Defaults to 14.
        in_channels (int, optional): Number of input channels. Defaults to 3.
        hidden_size (int, optional): Hidden dimension. Defaults to 1024 (Large).
        num_layers (int, optional): Number of transformer layers. Defaults to 24 (Large).
        num_heads (int, optional): Number of attention heads. Defaults to 8 (Large).
        mlp_dim (int, optional): MLP intermediate dimension. Defaults to 2816 (Large).
        rms_norm_eps (float, optional): RMSNorm epsilon. Defaults to 1e-5.
        use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
        attention_fn (Callable[..., Any] | None, optional): Custom attention function. Defaults to None.
        rngs (rnglib.Rngs | None, optional): RNG state. If None, initializes to nnx.Rngs(0).
        dtype (DTypeLike, optional): Computation dtype. Defaults to jnp.float32.
        param_dtype (DTypeLike, optional): Parameter dtype. Defaults to jnp.float32.
        sharding (ShardingSpec | None, optional): Sharding specification. Defaults to AIMv2Sharding().
    """
    if rngs is None:
        rngs = nnx.Rngs(0)
    if sharding is None:
        sharding = AIMv2Sharding()
    self._original_config = None
    self.encoder = VisionTransformerBase(
        img_size=img_size,
        patch_size=patch_size,
        in_channels=in_channels,
        hidden_size=hidden_size,
        num_layers=num_layers,
        num_heads=num_heads,
        mlp_dim=mlp_dim,
        layernorm_epsilon=rms_norm_eps,
        pooling_type="ALL",
        dropout_rate=0.0,
        act_fn=jax.nn.silu,
        use_pre_norm=True,
        pre_norm_before_pos=True,
        use_patch_bias=True,
        use_gradient_checkpointing=use_gradient_checkpointing,
        use_layer_scale=False,
        use_rope=False,
        num_register_tokens=0,
        use_gated_mlp=True,
        key_bias=False,
        attn_bias=False,
        mlp_bias=False,
        use_rms_norm=True,
        attention_fn=attention_fn,
        rngs=rngs,
        dtype=dtype,
        param_dtype=param_dtype,
        sharding=sharding,
    )

from_config(config, *, rngs=None, dtype=jnp.float32, param_dtype=jnp.float32, sharding=None, use_gradient_checkpointing=False, attention_fn=None) classmethod

Create AIMv2Model from a HuggingFace-compatible config dict.

Parameters:

Name Type Description Default
config dict[str, Any]

Configuration dictionary in HuggingFace AIMv2 format.

 required
rngs Rngs | None

RNG state. If None, initializes to nnx.Rngs(0).

 None
dtype DTypeLike

Computation dtype. Defaults to jnp.float32.

 float32
param_dtype DTypeLike

Parameter dtype. Defaults to jnp.float32.

 float32
sharding ShardingSpec | None

Sharding specification. Defaults to AIMv2Sharding().

 None
use_gradient_checkpointing bool

Whether to use gradient checkpointing. Defaults to False.

 False
attention_fn Callable[..., Any] | None

Custom attention function. Defaults to None.

 None

Returns:

Name Type Description
AIMv2Model AIMv2Model

Model with randomly initialized weights matching the given config.
Source code in src/jimm/models/aimv2/aimv2_model.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
@classmethod
def from_config(
    cls,
    config: dict[str, Any],
    *,
    rngs: rnglib.Rngs | None = None,
    dtype: DTypeLike = jnp.float32,
    param_dtype: DTypeLike = jnp.float32,
    sharding: ShardingSpec | None = None,
    use_gradient_checkpointing: bool = False,
    attention_fn: Callable[..., Any] | None = None,
) -> "AIMv2Model":
    """Create AIMv2Model from a HuggingFace-compatible config dict.

    Args:
        config (dict[str, Any]): Configuration dictionary in HuggingFace AIMv2 format.
        rngs (rnglib.Rngs | None): RNG state. If None, initializes to nnx.Rngs(0).
        dtype (DTypeLike): Computation dtype. Defaults to jnp.float32.
        param_dtype (DTypeLike): Parameter dtype. Defaults to jnp.float32.
        sharding (ShardingSpec | None): Sharding specification. Defaults to AIMv2Sharding().
        use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
        attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

    Returns:
        AIMv2Model: Model with randomly initialized weights matching the given config.
    """
    if rngs is None:
        rngs = nnx.Rngs(0)
    return cls(
        **cls._parse_config(config),
        use_gradient_checkpointing=use_gradient_checkpointing,
        attention_fn=attention_fn,
        rngs=rngs,
        dtype=dtype,
        param_dtype=param_dtype,
        sharding=sharding,
    )

from_pretrained(model_name_or_path, use_pytorch=False, rngs=None, dtype=jnp.float32, param_dtype=jnp.float32, sharding=None, use_gradient_checkpointing=False, attention_fn=None) classmethod

Load a pretrained AIMv2 model from a local path or HuggingFace Hub.

Parameters:

Name Type Description Default
model_name_or_path str

Local directory or HuggingFace model ID (e.g. "apple/aimv2-large-patch14-224").

 required
use_pytorch bool

Whether to load from PyTorch weights. Defaults to False.

 False
rngs Rngs | None

RNG state. If None, initializes to nnx.Rngs(0).

 None
dtype DTypeLike

Computation dtype. Defaults to jnp.float32.

 float32
param_dtype DTypeLike

Parameter dtype. Defaults to jnp.float32.

 float32
sharding ShardingSpec | None

Sharding specification. Defaults to AIMv2Sharding().

 None
use_gradient_checkpointing bool

Whether to use gradient checkpointing. Defaults to False.

 False
attention_fn Callable[..., Any] | None

Custom attention function. Defaults to None.

 None

Returns:

Name Type Description
AIMv2Model AIMv2Model

Model with pretrained weights loaded.
Source code in src/jimm/models/aimv2/aimv2_model.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
@classmethod
def from_pretrained(
    cls,
    model_name_or_path: str,
    use_pytorch: bool = False,
    rngs: rnglib.Rngs | None = None,
    dtype: DTypeLike = jnp.float32,
    param_dtype: DTypeLike = jnp.float32,
    sharding: ShardingSpec | None = None,
    use_gradient_checkpointing: bool = False,
    attention_fn: Callable[..., Any] | None = None,
) -> "AIMv2Model":
    """Load a pretrained AIMv2 model from a local path or HuggingFace Hub.

    Args:
        model_name_or_path (str): Local directory or HuggingFace model ID
            (e.g. "apple/aimv2-large-patch14-224").
        use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
        rngs (rnglib.Rngs | None): RNG state. If None, initializes to nnx.Rngs(0).
        dtype (DTypeLike): Computation dtype. Defaults to jnp.float32.
        param_dtype (DTypeLike): Parameter dtype. Defaults to jnp.float32.
        sharding (ShardingSpec | None): Sharding specification. Defaults to AIMv2Sharding().
        use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
        attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

    Returns:
        AIMv2Model: Model with pretrained weights loaded.
    """
    from .params import load_from_pretrained

    return load_from_pretrained(
        cls,
        model_name_or_path,
        use_pytorch,
        rngs,
        dtype,
        param_dtype,
        sharding,
        use_gradient_checkpointing,
        attention_fn,
    )

save_pretrained(save_directory)

Save model weights and config in HuggingFace format.

Parameters:

Name Type Description Default
save_directory str

Directory path where the model will be saved.

 required
Source code in src/jimm/models/aimv2/aimv2_model.py 
107
108
109
110
111
112
113
114
115
def save_pretrained(self, save_directory: str) -> None:
    """Save model weights and config in HuggingFace format.

    Args:
        save_directory (str): Directory path where the model will be saved.
    """
    from .params import save_pretrained as _save_pretrained

    _save_pretrained(self, save_directory)