Skip to content

DINOv3

DINOv3 is a self-supervised vision transformer that extends the DINOv2 training recipe with three architectural improvements: 2D Rotary Position Embeddings (RoPE), register tokens, and a gated MLP (SwiGLU). It is trained on LVD-1689M, a curated dataset of 1.69 billion images, and produces general-purpose image representations that transfer well across downstream tasks.

Paper: "DINOv3" (Siméoni et al., 2025) Code: github.com/facebookresearch/dinov3

The jimm implementation returns the CLS token embedding after the final LayerNorm. Because RoPE computes position embeddings dynamically from the input spatial dimensions, the model accepts variable image sizes: any height and width divisible by patch_size, without retraining or interpolation.

Key architectural differences from DINOv2:

Feature DINOv2 DINOv3
Position embeddings Learned absolute (fixed grid) 2D RoPE (dynamic, variable size)
Register tokens None 4 learnable tokens between CLS and patches
MLP Standard (fc1 -> GELU -> fc2) Gated SwiGLU
Key bias Yes No

Register tokens were introduced in "Vision Transformers Need Registers" (Darcet et al., 2024).

Supported models

HuggingFace ID hidden_size Layers Heads Patch Image size
facebook/dinov3-vits16-pretrain-lvd1689m 384 12 6 16 variable
facebook/dinov3-vitb16-pretrain-lvd1689m 768 12 12 16 variable
facebook/dinov3-vitl16-pretrain-lvd1689m 1024 24 16 16 variable

Note: These are gated-repo models. You must request access on HuggingFace and authenticate with huggingface-cli login (or set HF_TOKEN) before loading.

Basic usage

import jimm
import numpy as np

model = jimm.DINOv3Model.from_pretrained("facebook/dinov3-vits16-pretrain-lvd1689m")

# Standard 224x224 input
images = np.random.rand(4, 224, 224, 3).astype(np.float32)
embeddings = model(images)  # shape: (4, 384)

# Variable image size -- any multiple of patch_size (16)
images_rect = np.random.rand(4, 192, 256, 3).astype(np.float32)
embeddings_rect = model(images_rect)  # shape: (4, 384)

Variable image sizes

Unlike DINOv2's fixed learned position embedding table, DINOv3's RoPE is recomputed on every forward pass from the actual spatial dimensions. Any height and width divisible by patch_size work without interpolation:

import jimm

model = jimm.DINOv3Model.from_pretrained("facebook/dinov3-vits16-pretrain-lvd1689m")

model(images_224)      # 224x224 -> 196 patch tokens
model(images_336)      # 336x336 -> 441 patch tokens
model(images_192x256)  # 192x256 -> 192 patch tokens

Note: JAX traces the model once per unique (height, width) pair. For production use, fix the image size or pre-warm the shapes you need.

Flash / Splash Attention

DINOv3 does not support custom attention functions. The RoPE path applies Q/K rotations and scaled dot-product attention directly in JAX, bypassing nnx.MultiHeadAttention.__call__. Passing attention_fn raises ValueError.

FSDP / Explicit Sharding

DINOv3 supports JAX explicit sharding (FSDP-style) via DINOv3Sharding, which extends ViTSharding with sharding for LayerScale vectors and gated MLP biases.

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.DINOv3Model.from_pretrained("facebook/dinov3-vits16-pretrain-lvd1689m")
# model params are automatically sharded across fsdp axis

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

jimm.models.dinov3.DINOv3Model

Bases: Module

DINOv3 vision transformer feature extractor.

Implements the DINOv3 architecture with 2D rotary position embeddings (RoPE), register tokens, and optional gated MLP. Returns CLS token embeddings. Supports variable image sizes (any multiple of patch_size) via dynamic RoPE.

Source code in src/jimm/models/dinov3/dinov3_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
 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
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
class DINOv3Model(nnx.Module):
    """DINOv3 vision transformer feature extractor.

    Implements the DINOv3 architecture with 2D rotary position embeddings (RoPE),
    register tokens, and optional gated MLP. Returns CLS token embeddings.
    Supports variable image sizes (any multiple of patch_size) via dynamic RoPE.
    """

    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 16,
        in_channels: int = 3,
        hidden_size: int = 384,
        num_layers: int = 12,
        num_heads: int = 6,
        mlp_dim: int = 1536,
        num_register_tokens: int = 4,
        rope_theta: float = 100.0,
        layer_scale_init: float = 1.0,
        layernorm_epsilon: float = 1e-5,
        act_fn: Callable | None = None,
        use_gated_mlp: bool = False,
        use_patch_bias: bool = True,
        key_bias: bool = False,
        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 = DINOv3Sharding(),
    ) -> None:
        """Initialize DINOv3Model.

        Args:
            img_size (int, optional): Input image size hint (ignored at runtime — variable size supported). Defaults to 224.
            patch_size (int, optional): Patch size. Defaults to 16.
            in_channels (int, optional): Number of input channels. Defaults to 3.
            hidden_size (int, optional): Hidden dimension size. Defaults to 384 (small).
            num_layers (int, optional): Number of transformer layers. Defaults to 12.
            num_heads (int, optional): Number of attention heads. Defaults to 6.
            mlp_dim (int, optional): MLP intermediate dimension. Defaults to 1536.
            num_register_tokens (int, optional): Number of learnable register tokens. Defaults to 4.
            rope_theta (float, optional): RoPE base frequency. Defaults to 100.0.
            layer_scale_init (float, optional): Initial value for per-channel LayerScale parameters. Defaults to 1.0.
            layernorm_epsilon (float, optional): LayerNorm epsilon. Defaults to 1e-5.
            act_fn (Callable | None, optional): MLP activation function. Defaults to exact GELU.
            use_gated_mlp (bool, optional): Whether to use gated (SwiGLU-style) MLP. Defaults to False.
                Note: HuggingFace DINOv3 checkpoints typically use True; use from_pretrained or from_config
                to load the correct architecture automatically rather than relying on this default.
            use_patch_bias (bool, optional): Whether to include a bias in the patch embedding Conv. Defaults to True.
            key_bias (bool, optional): Whether to include a bias in the key projection. Defaults to False.
            use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
            attention_fn (Callable[..., Any] | None, optional): Must be None. DINOv3 uses a manual RoPE attention
                path that bypasses nnx.MultiHeadAttention.__call__; passing a custom attention_fn raises ValueError.
            rngs (rnglib.Rngs | None, optional): The random number generator state. If None, initializes to nnx.Rngs(0).
            dtype (DTypeLike, optional): The data type for computations. Defaults to jnp.float32.
            param_dtype (DTypeLike, optional): The data type for parameters. Defaults to jnp.float32.
            sharding (ShardingSpec, optional): Sharding specification for parameters. Defaults to DINOv3Sharding.
        """
        if rngs is None:
            rngs = nnx.Rngs(0)
        if attention_fn is not None:
            raise ValueError("DINOv3 uses a manual RoPE attention path that bypasses nnx.MultiHeadAttention.__call__; attention_fn is not supported. Remove the attention_fn argument.")
        _act_fn = act_fn if act_fn is not None else _DEFAULT_ACT_FN
        self._original_config = None
        self._layer_scale_init = layer_scale_init
        self._act_fn = _act_fn
        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,
            pooling_type="CLS",
            dropout_rate=0.0,
            use_pre_norm=False,
            use_patch_bias=use_patch_bias,
            layernorm_epsilon=layernorm_epsilon,
            use_layer_scale=True,
            layer_scale_init=layer_scale_init,
            use_rope=True,
            rope_theta=rope_theta,
            num_register_tokens=num_register_tokens,
            use_gated_mlp=use_gated_mlp,
            act_fn=_act_fn,
            key_bias=key_bias,
            use_gradient_checkpointing=use_gradient_checkpointing,
            attention_fn=None,
            rngs=rngs,
            dtype=dtype,
            param_dtype=param_dtype,
            sharding=sharding,
        )

    def __call__(self, x: Float[Array, "batch height width channels"]) -> Float[Array, "batch hidden_size"]:
        """Run DINOv3 forward pass and return CLS token embedding.

        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 hidden_size"]: CLS token embeddings after final LayerNorm.
        """
        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 = DINOv3Sharding(),
        use_gradient_checkpointing: bool = False,
        attention_fn: Callable[..., Any] | None = None,
    ) -> "DINOv3Model":
        """Load a pretrained DINOv3 model from a local path or HuggingFace Hub.

        Args:
            model_name_or_path (str): Local directory or HuggingFace model ID (e.g. "facebook/dinov3-vits16-pretrain-lvd1689m").
            use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
            rngs (rnglib.Rngs | None): The random number generator state. If None, initializes to nnx.Rngs(0).
            dtype (DTypeLike): The data type for computations. Defaults to jnp.float32.
            param_dtype (DTypeLike): The data type for parameters. Defaults to jnp.float32.
            sharding (ShardingSpec): Sharding specification for parameters. Defaults to DINOv3Sharding.
            use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
            attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

        Returns:
            DINOv3Model: 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]:
        hidden_size = config["hidden_size"]
        mlp_dim = config.get("intermediate_size", int(hidden_size * config.get("mlp_ratio", 4)))
        hidden_act_str = config.get("hidden_act", "gelu")
        act_fn = _ACT_FN_MAP.get(hidden_act_str, _DEFAULT_ACT_FN)
        return {
            "img_size": config.get("image_size", 224),
            "patch_size": config["patch_size"],
            "in_channels": config.get("num_channels", 3),
            "hidden_size": hidden_size,
            "num_layers": config["num_hidden_layers"],
            "num_heads": config["num_attention_heads"],
            "mlp_dim": mlp_dim,
            "num_register_tokens": config.get("num_register_tokens", 4),
            "rope_theta": config.get("rope_theta", 100.0),
            "layer_scale_init": config.get("layerscale_value", 1.0),
            "layernorm_epsilon": config.get("layer_norm_eps", 1e-5),
            "act_fn": act_fn,
            "use_gated_mlp": config.get("use_gated_mlp", False),
            "use_patch_bias": config.get("use_patch_bias", True),
            "key_bias": config.get("key_bias", False),
        }

    @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 = DINOv3Sharding(),
        use_gradient_checkpointing: bool = False,
        attention_fn: Callable[..., Any] | None = None,
    ) -> "DINOv3Model":
        """Create DINOv3Model from a HuggingFace-compatible config dict.

        Args:
            config (dict[str, Any]): Configuration dictionary in HuggingFace DINOv3 format.
            rngs (rnglib.Rngs | None): The random number generator state. If None, initializes to nnx.Rngs(0).
            dtype (DTypeLike): The data type for computations. Defaults to jnp.float32.
            param_dtype (DTypeLike): The data type for parameters. Defaults to jnp.float32.
            sharding (ShardingSpec): Sharding specification for parameters. Defaults to DINOv3Sharding.
            use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
            attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

        Returns:
            DINOv3Model: 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 DINOv3 forward pass and return CLS token embedding.

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 hidden_size']

Float[Array, "batch hidden_size"]: CLS token embeddings after final LayerNorm.
Source code in src/jimm/models/dinov3/dinov3_model.py 
121
122
123
124
125
126
127
128
129
130
131
def __call__(self, x: Float[Array, "batch height width channels"]) -> Float[Array, "batch hidden_size"]:
    """Run DINOv3 forward pass and return CLS token embedding.

    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 hidden_size"]: CLS token embeddings after final LayerNorm.
    """
    return self.encoder(x)

__init__(img_size=224, patch_size=16, in_channels=3, hidden_size=384, num_layers=12, num_heads=6, mlp_dim=1536, num_register_tokens=4, rope_theta=100.0, layer_scale_init=1.0, layernorm_epsilon=1e-05, act_fn=None, use_gated_mlp=False, use_patch_bias=True, key_bias=False, use_gradient_checkpointing=False, attention_fn=None, rngs=None, dtype=jnp.float32, param_dtype=jnp.float32, sharding=DINOv3Sharding())

Initialize DINOv3Model.

Parameters:

Name Type Description Default
img_size int

Input image size hint (ignored at runtime — variable size supported). Defaults to 224.

 224
patch_size int

Patch size. Defaults to 16.

 16
in_channels int

Number of input channels. Defaults to 3.

 3
hidden_size int

Hidden dimension size. Defaults to 384 (small).

 384
num_layers int

Number of transformer layers. Defaults to 12.

 12
num_heads int

Number of attention heads. Defaults to 6.

 6
mlp_dim int

MLP intermediate dimension. Defaults to 1536.

 1536
num_register_tokens int

Number of learnable register tokens. Defaults to 4.

 4
rope_theta float

RoPE base frequency. Defaults to 100.0.

 100.0
layer_scale_init float

Initial value for per-channel LayerScale parameters. Defaults to 1.0.

 1.0
layernorm_epsilon float

LayerNorm epsilon. Defaults to 1e-5.

 1e-05
act_fn Callable | None

MLP activation function. Defaults to exact GELU.

 None
use_gated_mlp bool

Whether to use gated (SwiGLU-style) MLP. Defaults to False. Note: HuggingFace DINOv3 checkpoints typically use True; use from_pretrained or from_config to load the correct architecture automatically rather than relying on this default.

 False
use_patch_bias bool

Whether to include a bias in the patch embedding Conv. Defaults to True.

 True
key_bias bool

Whether to include a bias in the key projection. Defaults to False.

 False
use_gradient_checkpointing bool

Whether to use gradient checkpointing. Defaults to False.

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

Must be None. DINOv3 uses a manual RoPE attention path that bypasses nnx.MultiHeadAttention.call; passing a custom attention_fn raises ValueError.

 None
rngs Rngs | None

The random number generator state. If None, initializes to nnx.Rngs(0).

 None
dtype DTypeLike

The data type for computations. Defaults to jnp.float32.

 float32
param_dtype DTypeLike

The data type for parameters. Defaults to jnp.float32.

 float32
sharding ShardingSpec

Sharding specification for parameters. Defaults to DINOv3Sharding.

 DINOv3Sharding()
Source code in src/jimm/models/dinov3/dinov3_model.py 
 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
def __init__(
    self,
    img_size: int = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    hidden_size: int = 384,
    num_layers: int = 12,
    num_heads: int = 6,
    mlp_dim: int = 1536,
    num_register_tokens: int = 4,
    rope_theta: float = 100.0,
    layer_scale_init: float = 1.0,
    layernorm_epsilon: float = 1e-5,
    act_fn: Callable | None = None,
    use_gated_mlp: bool = False,
    use_patch_bias: bool = True,
    key_bias: bool = False,
    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 = DINOv3Sharding(),
) -> None:
    """Initialize DINOv3Model.

    Args:
        img_size (int, optional): Input image size hint (ignored at runtime — variable size supported). Defaults to 224.
        patch_size (int, optional): Patch size. Defaults to 16.
        in_channels (int, optional): Number of input channels. Defaults to 3.
        hidden_size (int, optional): Hidden dimension size. Defaults to 384 (small).
        num_layers (int, optional): Number of transformer layers. Defaults to 12.
        num_heads (int, optional): Number of attention heads. Defaults to 6.
        mlp_dim (int, optional): MLP intermediate dimension. Defaults to 1536.
        num_register_tokens (int, optional): Number of learnable register tokens. Defaults to 4.
        rope_theta (float, optional): RoPE base frequency. Defaults to 100.0.
        layer_scale_init (float, optional): Initial value for per-channel LayerScale parameters. Defaults to 1.0.
        layernorm_epsilon (float, optional): LayerNorm epsilon. Defaults to 1e-5.
        act_fn (Callable | None, optional): MLP activation function. Defaults to exact GELU.
        use_gated_mlp (bool, optional): Whether to use gated (SwiGLU-style) MLP. Defaults to False.
            Note: HuggingFace DINOv3 checkpoints typically use True; use from_pretrained or from_config
            to load the correct architecture automatically rather than relying on this default.
        use_patch_bias (bool, optional): Whether to include a bias in the patch embedding Conv. Defaults to True.
        key_bias (bool, optional): Whether to include a bias in the key projection. Defaults to False.
        use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
        attention_fn (Callable[..., Any] | None, optional): Must be None. DINOv3 uses a manual RoPE attention
            path that bypasses nnx.MultiHeadAttention.__call__; passing a custom attention_fn raises ValueError.
        rngs (rnglib.Rngs | None, optional): The random number generator state. If None, initializes to nnx.Rngs(0).
        dtype (DTypeLike, optional): The data type for computations. Defaults to jnp.float32.
        param_dtype (DTypeLike, optional): The data type for parameters. Defaults to jnp.float32.
        sharding (ShardingSpec, optional): Sharding specification for parameters. Defaults to DINOv3Sharding.
    """
    if rngs is None:
        rngs = nnx.Rngs(0)
    if attention_fn is not None:
        raise ValueError("DINOv3 uses a manual RoPE attention path that bypasses nnx.MultiHeadAttention.__call__; attention_fn is not supported. Remove the attention_fn argument.")
    _act_fn = act_fn if act_fn is not None else _DEFAULT_ACT_FN
    self._original_config = None
    self._layer_scale_init = layer_scale_init
    self._act_fn = _act_fn
    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,
        pooling_type="CLS",
        dropout_rate=0.0,
        use_pre_norm=False,
        use_patch_bias=use_patch_bias,
        layernorm_epsilon=layernorm_epsilon,
        use_layer_scale=True,
        layer_scale_init=layer_scale_init,
        use_rope=True,
        rope_theta=rope_theta,
        num_register_tokens=num_register_tokens,
        use_gated_mlp=use_gated_mlp,
        act_fn=_act_fn,
        key_bias=key_bias,
        use_gradient_checkpointing=use_gradient_checkpointing,
        attention_fn=None,
        rngs=rngs,
        dtype=dtype,
        param_dtype=param_dtype,
        sharding=sharding,
    )

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

Create DINOv3Model from a HuggingFace-compatible config dict.

Parameters:

Name Type Description Default
config dict[str, Any]

Configuration dictionary in HuggingFace DINOv3 format.

 required
rngs Rngs | None

The random number generator state. If None, initializes to nnx.Rngs(0).

 None
dtype DTypeLike

The data type for computations. Defaults to jnp.float32.

 float32
param_dtype DTypeLike

The data type for parameters. Defaults to jnp.float32.

 float32
sharding ShardingSpec

Sharding specification for parameters. Defaults to DINOv3Sharding.

 DINOv3Sharding()
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
DINOv3Model DINOv3Model

Model with randomly initialized weights matching the given config.
Source code in src/jimm/models/dinov3/dinov3_model.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
@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 = DINOv3Sharding(),
    use_gradient_checkpointing: bool = False,
    attention_fn: Callable[..., Any] | None = None,
) -> "DINOv3Model":
    """Create DINOv3Model from a HuggingFace-compatible config dict.

    Args:
        config (dict[str, Any]): Configuration dictionary in HuggingFace DINOv3 format.
        rngs (rnglib.Rngs | None): The random number generator state. If None, initializes to nnx.Rngs(0).
        dtype (DTypeLike): The data type for computations. Defaults to jnp.float32.
        param_dtype (DTypeLike): The data type for parameters. Defaults to jnp.float32.
        sharding (ShardingSpec): Sharding specification for parameters. Defaults to DINOv3Sharding.
        use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
        attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

    Returns:
        DINOv3Model: 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=DINOv3Sharding(), use_gradient_checkpointing=False, attention_fn=None) classmethod

Load a pretrained DINOv3 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. "facebook/dinov3-vits16-pretrain-lvd1689m").

 required
use_pytorch bool

Whether to load from PyTorch weights. Defaults to False.

 False
rngs Rngs | None

The random number generator state. If None, initializes to nnx.Rngs(0).

 None
dtype DTypeLike

The data type for computations. Defaults to jnp.float32.

 float32
param_dtype DTypeLike

The data type for parameters. Defaults to jnp.float32.

 float32
sharding ShardingSpec

Sharding specification for parameters. Defaults to DINOv3Sharding.

 DINOv3Sharding()
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
DINOv3Model DINOv3Model

Model with pretrained weights loaded.
Source code in src/jimm/models/dinov3/dinov3_model.py 
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
@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 = DINOv3Sharding(),
    use_gradient_checkpointing: bool = False,
    attention_fn: Callable[..., Any] | None = None,
) -> "DINOv3Model":
    """Load a pretrained DINOv3 model from a local path or HuggingFace Hub.

    Args:
        model_name_or_path (str): Local directory or HuggingFace model ID (e.g. "facebook/dinov3-vits16-pretrain-lvd1689m").
        use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
        rngs (rnglib.Rngs | None): The random number generator state. If None, initializes to nnx.Rngs(0).
        dtype (DTypeLike): The data type for computations. Defaults to jnp.float32.
        param_dtype (DTypeLike): The data type for parameters. Defaults to jnp.float32.
        sharding (ShardingSpec): Sharding specification for parameters. Defaults to DINOv3Sharding.
        use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
        attention_fn (Callable[..., Any] | None): Custom attention function. Defaults to None.

    Returns:
        DINOv3Model: 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/dinov3/dinov3_model.py 
133
134
135
136
137
138
139
140
141
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)