jax-image-models 0.3.3__py3-none-any.whl

jax_image_models-0.3.3.dist-info/METADATA

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: jax-image-models
+Version: 0.3.3
+Summary: Jax Image Modeling of Models
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: flax>=0.10.6
+Requires-Dist: jax>=0.6.2
+Requires-Dist: jaxtyping>=0.3.2
+Requires-Dist: safetensors>=0.5.3
+Requires-Dist: tokamax>=0.0.12
+Description-Content-Type: text/markdown
+
+# Jax Image Modeling of Models (jimm)
+Docs are at: [https://pythoncrazy.github.io/jimm](https://pythoncrazy.github.io/jimm)
+- This aims to be the jax counterpart to timm, with the exception that for image-text models (CLIP, SigLIP, etc), we support the text model entirely.
+- Made with flax nnx, supports weight loading from pytorch_model.bin and safetensors (as well as both methods from huggingface).
+
+Models Supported:
+- Vision Transformers
+    - Both with a classification linear layer, or not
+    - Using a CLS Token for pooling, or using Multihead Attention Pooling
+    - Can load any standard variant of Vision Transformers of any size/resolution(e.g. "google/vit-base-patch16-224" or "google/vit-large-patch16-384")
+- CLIP
+    - Can load from any checkpoints of the clip model on github (such as "openai/clip-vit-base-patch32" or "geolocal/StreetCLIP")
+- SigLIP
+    - Can load any non-naflex version of the SigLIP model, from both siglipv1 and siglipv2 (eg "google/siglip-base-patch16-256" or "google/siglip2-large-patch16-512" from huggingface or locally)
+## Installation
+### Using pixi.sh:
+`pixi add jimm@https://github.com/pythoncrazy/jimm.git --pypi`
+### Using uv
+`uv add --dev git+https://github.com/pythoncrazy/jimm.git`
+or if you prefer to not add as a direct dependency:
+`uv pip install git+https://github.com/pythoncrazy/jimm.git`
+### Using pip/conda
+`pip install git+https://github.com/pythoncrazy/jimm.git`

jax_image_models-0.3.3.dist-info/RECORD

@@ -0,0 +1,25 @@
+jimm/__init__.py,sha256=7rsfJf4ma5htZG7IaXCklJ52m9fELLZiLNyxqdtdkBk,849
+jimm/common/autotuning.py,sha256=8E-S2gtHwSpJKZgR-BQlsosebJH_zn195HVeDhJClIE,4998
+jimm/common/loading_utils.py,sha256=zxwytlhLD26dpWrnCjd4LAQmLtWeD9Z2GSDH-6qt4Rc,9398
+jimm/common/sharding.py,sha256=_hXhuX4C6lneNHo28gK7CkBDxlvPujzlLlG65MzgVjs,4161
+jimm/common/tokamax_attention.py,sha256=CSJs_s2XKniuq8__fjTUoPoprMsqweNCpUQmwr_whQE,4160
+jimm/common/transformer.py,sha256=h-cpzg8Os-nU39izuh6bt3UQBwlsVK21DTh0BOie5m4,12194
+jimm/common/utils.py,sha256=r6sbZVmeZbwj6opK6MTf61Y7w4d_3zmuXVFkWQoSXcE,6706
+jimm/common/vit.py,sha256=_pZv4hXk67Xg0hwjkQnoQ_Q5RXwyGO1Hw7hokdfWLm4,14940
+jimm/models/__init__.py,sha256=dkC9Ujsrr0qEjPNl2vaoTJEOOADy766ZOUEKVsCbrbQ,311
+jimm/models/clip/__init__.py,sha256=Ayy0m9sNHurvQy3l-OlhWoc_NsGeCUN8Oxp_Jw7NHbg,117
+jimm/models/clip/clip_model.py,sha256=K3oyN1yrisrOf1JpspiLPOD52j-ieeSTPub0sEX4EQs,27822
+jimm/models/clip/params.py,sha256=mZOGjL0at_uv3dljiGbIumkSkXjf-J0zf958vTuwUKo,24504
+jimm/models/clip/sharding.py,sha256=KIfUU_27dJTk75r1QZtfaBzmvd2E8udS-4gzp4w2jxc,3115
+jimm/models/siglip/__init__.py,sha256=2-hoGa2kAYhsNC-n52m_5PTE35XYNGXuFAQSnMtQ2QE,131
+jimm/models/siglip/params.py,sha256=sc4_x8332boIXNg_FpnfIQKkFnBUpmPMuIcycZudOJs,29485
+jimm/models/siglip/sharding.py,sha256=Ck0HRAYz4gYN5XJws05lKxad04GfpylszaSM2KMN0zc,3029
+jimm/models/siglip/siglip_model.py,sha256=k-yf92I18o28rkBBO3EhGdW-zNiGY39B7hVCQbhdQbQ,26495
+jimm/models/vit/__init__.py,sha256=Hw1rwR3mN13X-0TYmq-H2cZsXUmK95326sYADF3VjWM,74
+jimm/models/vit/params.py,sha256=IaP0Z4nzdXIXWm28P-x8ERJ6lNNrd8HbL7CXFxigYd0,12890
+jimm/models/vit/sharding.py,sha256=8LLOvAOfY0770jF3GZth4JiWzuqMZfERfMqdJ7a-vsg,2963
+jimm/models/vit/vit_model.py,sha256=WenSypHga0G_4kCyBUR5Gru5wDFsVNcEqst3rl6B1qo,9057
+jax_image_models-0.3.3.dist-info/METADATA,sha256=PfaDFPWkrf1BZtIEJGDWd0Gtx2TnSNnGH7Vk5w_teoA,1762
+jax_image_models-0.3.3.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+jax_image_models-0.3.3.dist-info/licenses/LICENSE,sha256=wYMBl9KN_WLSFwJ2lMHMchvv7IhGvIPl4alCEzjFw6I,1070
+jax_image_models-0.3.3.dist-info/RECORD,,

jax_image_models-0.3.3.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

jax_image_models-0.3.3.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Pinak Paliwal
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jimm/__init__.py

@@ -0,0 +1,33 @@
+from .common.autotuning import AutotuningResult
+from .common.autotuning import autotune
+from .common.autotuning import autotuned_fn
+from .common.autotuning import cached_autotune
+from .common.autotuning import load_autotune_result
+from .common.tokamax_attention import make_tokamax_attention
+from .common.tokamax_attention import tokamax_attention_fn as tokamax_attention
+from .models import (
+    CLIP,
+    CLIPTextModel,
+    CLIPVisionModel,
+    SigLIP,
+    SigLIPTextModel,
+    SigLIPVisionModel,
+    VisionTransformer,
+)
+
+__all__ = [
+    "tokamax_attention",
+    "make_tokamax_attention",
+    "autotune",
+    "cached_autotune",
+    "autotuned_fn",
+    "load_autotune_result",
+    "AutotuningResult",
+    "VisionTransformer",
+    "CLIP",
+    "CLIPTextModel",
+    "CLIPVisionModel",
+    "SigLIP",
+    "SigLIPTextModel",
+    "SigLIPVisionModel",
+]

jimm/common/autotuning.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import contextlib
+import functools
+import hashlib
+import os
+from collections.abc import Callable
+from typing import Any
+
+from jax.extend import backend as _jax_backend
+from tokamax._src.autotuning import api as _autotune_api
+
+AutotuningResult = _autotune_api.AutotuningResult
+
+
+def _lower(jitted_fn: Any, *args: Any) -> Any:
+    """Unwrap a Flax ``nnx.jit`` ``Lowered`` to the raw ``jax.stages.Lowered``.
+
+    ``nnx.jit`` wraps the JAX lowered object; tokamax's HLO utils require the
+    raw type, accessible via the ``.lowered`` attribute.
+
+    Args:
+        jitted_fn (Any): An ``nnx.jit`` or ``jax.jit`` callable.
+        *args (Any): Sample arguments for lowering.
+
+    Returns:
+        Any: Raw ``jax.stages.Lowered`` computation.
+    """
+    lowered = jitted_fn.lower(*args)
+    return getattr(lowered, "lowered", lowered)
+
+
+def autotune(
+    jitted_fn: Any,
+    *sample_args: Any,
+    save_path: str | os.PathLike | None = None,
+    **kwargs: Any,
+) -> AutotuningResult:
+    """Autotune all tokamax ops in a jitted function and optionally save results.
+
+    Lowers ``jitted_fn`` with ``sample_args`` to extract op shapes, then
+    microbenchmarks every kernel config. Use the result as a context manager::
+
+        result = jimm.autotune(forward, model, image, text)
+        with result:
+            out = forward(model, image, text)
+
+    Args:
+        jitted_fn (Any): Jitted callable containing tokamax ops.
+        *sample_args (Any): Representative inputs; shapes/dtypes must match
+            production inputs.
+        save_path (str | os.PathLike | None): Serialize result as JSON here.
+        **kwargs (Any): Forwarded to ``tokamax.autotune``
+            (e.g. ``all_implementations``, ``progress_bar``).
+
+    Returns:
+        AutotuningResult: Best config for each op found in ``jitted_fn``.
+    """
+    result = _autotune_api.autotune(_lower(jitted_fn, *sample_args), **kwargs)
+    if save_path is not None:
+        with open(save_path, "w") as f:
+            result.dump(f)
+    return result
+
+
+def load_autotune_result(path: str | os.PathLike) -> AutotuningResult:
+    """Load an :class:`AutotuningResult` from a JSON file.
+
+    Args:
+        path (str | os.PathLike): Path written by :func:`autotune` or
+            :func:`cached_autotune`.
+
+    Returns:
+        AutotuningResult: Deserialized result.
+    """
+    with open(path) as f:
+        return AutotuningResult.load(f)
+
+
+def cached_autotune(
+    jitted_fn: Any,
+    *sample_args: Any,
+    cache_dir: str | os.PathLike,
+    **kwargs: Any,
+) -> AutotuningResult | None:
+    """Load autotuning results from disk, or run autotune and cache them.
+
+    The cache key is an MD5 hash of op autotuning keys + device kind, so the
+    same entry is reused whenever model architecture and hardware are unchanged.
+
+    Args:
+        jitted_fn (Any): Jitted callable containing tokamax ops.
+        *sample_args (Any): Representative inputs used to extract op shapes.
+        cache_dir (str | os.PathLike): Directory for JSON cache files.
+        **kwargs (Any): Forwarded to ``tokamax.autotune``.
+
+    Returns:
+        AutotuningResult | None: Loaded or freshly computed result, or ``None``
+        if no tunable ops exist.
+    """
+    bound_args = _autotune_api.get_bound_args(_lower(jitted_fn, *sample_args))
+    if not bound_args:
+        return None
+
+    device_kind = _jax_backend.get_default_device().device_kind
+    key_parts = sorted(str(ba.autotuning_cache_key) for ba in bound_args)
+    cache_key = hashlib.md5((device_kind + ":" + "|".join(key_parts)).encode()).hexdigest()[:16]
+
+    os.makedirs(cache_dir, exist_ok=True)
+    cache_path = os.path.join(cache_dir, f"{cache_key}.json")
+
+    if os.path.exists(cache_path):
+        return load_autotune_result(cache_path)
+
+    result = _autotune_api.autotune(list(bound_args), **kwargs)
+    with open(cache_path, "w") as f:
+        result.dump(f)
+    return result
+
+
+def autotuned_fn(
+    jitted_fn: Any,
+    *sample_args: Any,
+    cache_dir: str | os.PathLike,
+    **kwargs: Any,
+) -> Callable[..., Any]:
+    """Wrap a jitted function to always run with tuned tokamax configs.
+
+    Eagerly calls :func:`cached_autotune`, then returns a wrapper that injects
+    the tuned configs via :class:`AutotuningResult` on every call.
+
+    Args:
+        jitted_fn (Any): Jitted callable containing tokamax ops.
+        *sample_args (Any): Representative inputs used for autotuning.
+        cache_dir (str | os.PathLike): Directory for the autotune cache.
+        **kwargs (Any): Forwarded to ``tokamax.autotune``.
+
+    Returns:
+        Callable[..., Any]: Wrapped callable that applies tuned configs on
+        every invocation.
+    """
+    result = cached_autotune(jitted_fn, *sample_args, cache_dir=cache_dir, **kwargs)
+    ctx = result if result is not None else contextlib.nullcontext()
+
+    @functools.wraps(jitted_fn)
+    def wrapper(*args: Any, **kw: Any) -> Any:
+        with ctx:
+            return jitted_fn(*args, **kw)
+
+    return wrapper

jimm/common/loading_utils.py

@@ -0,0 +1,248 @@
+import json
+import os
+import re
+from math import prod
+from typing import Any, Dict, Tuple, TypeVar
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from huggingface_hub import hf_hub_download
+from jaxtyping import Array, DTypeLike
+from safetensors.flax import load_file as load_safetensors_flax_file
+
+_M = TypeVar("_M", bound=nnx.Module)
+
+
+def load_params_and_config(
+    model_name_or_path: str,
+    use_pytorch: bool = False,
+    default_config_filename: str = "config.json",
+    default_pytorch_filename: str = "pytorch_model.bin",
+    default_safetensors_filename: str = "model.safetensors",
+) -> Tuple[Dict[str, Array], Dict[str, Any]]:
+    """Load model parameters and configuration from local directory or HuggingFace Hub.
+
+    Args:
+        model_name_or_path (str): Local directory path or HuggingFace model ID.
+        use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
+        default_config_filename (str): Config filename. Defaults to "config.json".
+        default_pytorch_filename (str): PyTorch weights filename. Defaults to "pytorch_model.bin".
+        default_safetensors_filename (str): Safetensors filename. Defaults to "model.safetensors".
+
+    Returns:
+        Tuple[Dict[str, Array], Dict[str, Any]]: Loaded parameters and configuration.
+    """
+    if os.path.isdir(model_name_or_path):
+        config_file_path = os.path.join(model_name_or_path, default_config_filename)
+        weights_filename = default_pytorch_filename if use_pytorch else default_safetensors_filename
+        weights_file_path = os.path.join(model_name_or_path, weights_filename)
+    else:
+        config_file_path = hf_hub_download(repo_id=model_name_or_path, filename=default_config_filename)
+        weights_filename = default_pytorch_filename if use_pytorch else default_safetensors_filename
+        weights_file_path = hf_hub_download(repo_id=model_name_or_path, filename=weights_filename)
+
+    with open(config_file_path, "r") as f:
+        config = json.load(f)
+
+    if use_pytorch:
+        import torch
+
+        state_dict = torch.load(weights_file_path, map_location="cpu")
+        params_fstate = {k: jnp.array(v.numpy()) for k, v in state_dict.items()}
+    else:
+        params_fstate = load_safetensors_flax_file(weights_file_path)
+
+    return params_fstate, config
+
+
+def stoi(k: str) -> int | str:
+    """Convert a string to int if numeric, else return as-is.
+
+    Args:
+        k (str): Key to convert.
+
+    Returns:
+        int | str: Integer if numeric string, else original string.
+    """
+    try:
+        return int(k)
+    except ValueError:
+        return k
+
+
+def map_to_bonsai_key(
+    mapping: dict[str, tuple[str, Any]],
+    key: str,
+) -> tuple[str | None, Any]:
+    """Match a flat HF key against regex patterns and return the flax target key.
+
+    Args:
+        mapping (dict[str, tuple[str, Any]]): Dict of {regex_pattern: (flax_key_template, transform)}.
+        key (str): HuggingFace parameter key to match.
+
+    Returns:
+        tuple[str | None, Any]: (flax_key, transform) if matched, else (None, None).
+    """
+    for pattern, (target, transform) in mapping.items():
+        if re.fullmatch(pattern, key):
+            return re.sub(pattern, target, key), transform
+    return None, None
+
+
+def to_scan_batched_keys(keys: tuple) -> tuple[tuple | None, int | None]:
+    """Convert a per-layer flat state key to its scan-batched equivalent.
+
+    With nnx.scan/vmap, transformer layers are stored as a single batched module
+    under a "layers" key rather than separate "layers_N" keys. This converts
+    keys like ("encoder", "layers_3", "attn", "query", "kernel") to
+    ("encoder", "layers", "attn", "query", "kernel") and returns the layer index.
+
+    Args:
+        keys (tuple): Flat state key tuple potentially containing a "layers_N" component.
+
+    Returns:
+        tuple[tuple | None, int | None]: (batched_keys, layer_idx) if a layers_N component
+            is found, else (None, None).
+    """
+    new_keys = list(keys)
+    layer_idx = None
+    for i, k in enumerate(new_keys):
+        if isinstance(k, str) and k.startswith("layers_"):
+            try:
+                layer_idx = int(k[7:])
+                new_keys[i] = "layers"
+                break
+            except ValueError:
+                pass
+    if layer_idx is None:
+        return None, None
+    return tuple(new_keys), layer_idx
+
+
+def _reshape_if_compatible(tensor: Array, target_shape: tuple[int, ...], hf_key: str, bonsai_keys: tuple[Any, ...]) -> Array:
+    """Reshape tensor only when element counts match, else raise a clear error."""
+    if tensor.shape == target_shape:
+        return tensor
+
+    if tensor.size != prod(target_shape):
+        bonsai_key = ".".join(str(key) for key in bonsai_keys)
+        raise ValueError(f"Shape mismatch for {hf_key} -> {bonsai_key}: got {tensor.shape}, expected {target_shape}")
+
+    return tensor.reshape(target_shape)
+
+
+def apply_mapping(
+    model: _M,
+    params_fstate: dict[str, Any],
+    mapping: dict[str, tuple[str, Any]],
+    param_dtype: DTypeLike,
+) -> _M:
+    """Apply regex-based HF parameter mappings to a model in-place."""
+    flat_state = dict(nnx.to_flat_state(nnx.state(model, nnx.Param)))
+    layer_accum: dict[tuple, dict[int, Any]] = {}
+
+    for hf_key, tensor in params_fstate.items():
+        bonsai_key, transform = map_to_bonsai_key(mapping, hf_key)
+        if bonsai_key is None:
+            continue
+
+        keys = tuple(stoi(k) for k in bonsai_key.split("."))
+        permute_rule, _, _ = transform.value
+        transformed = tensor.astype(param_dtype)
+        if permute_rule is not None:
+            transformed = jnp.transpose(transformed, permute_rule)
+
+        if keys in flat_state:
+            var = flat_state[keys]
+            var[...] = _reshape_if_compatible(transformed, var[...].shape, hf_key, keys)
+            continue
+
+        batched_keys, layer_idx = to_scan_batched_keys(keys)
+        if batched_keys is not None and layer_idx is not None and batched_keys in flat_state:
+            layer_accum.setdefault(batched_keys, {})[layer_idx] = transformed
+
+    for batched_keys, layers_dict in layer_accum.items():
+        var = flat_state[batched_keys]
+        num_layers = var[...].shape[0]
+        missing = sorted(set(range(num_layers)) - set(layers_dict))
+        if missing:
+            bonsai_key = ".".join(str(key) for key in batched_keys)
+            raise ValueError(f"Missing scanned layers for {bonsai_key}: {missing}")
+
+        stacked = jnp.stack([layers_dict[i] for i in range(num_layers)], axis=0)
+        var[...] = _reshape_if_compatible(stacked, var[...].shape, ".".join(str(key) for key in batched_keys), batched_keys)
+
+    nnx.update(model, nnx.from_flat_state(flat_state))
+    return model
+
+
+def _slice_layer(d: dict, idx: int) -> dict:
+    """Recursively extract index idx from batched arrays in a nested dict."""
+    result = {}
+    for key, value in d.items():
+        if isinstance(value, dict):
+            result[key] = _slice_layer(value, idx)
+        elif isinstance(value, jax.Array):
+            result[key] = value[idx]
+        else:
+            result[key] = value
+    return result
+
+
+def _infer_num_layers(d: dict) -> int | None:
+    """Return the leading dimension of the first JAX array found in d."""
+    for value in d.values():
+        if isinstance(value, jax.Array):
+            return int(value.shape[0])
+        if isinstance(value, dict):
+            n = _infer_num_layers(value)
+            if n is not None:
+                return n
+    return None
+
+
+def _is_numeric_key(key: Any) -> bool:
+    """Return True when a nested state-dict key is an integer layer index."""
+    return isinstance(key, int) or (isinstance(key, str) and key.isdigit())
+
+
+def _should_expand_layers_dict(d: dict) -> bool:
+    """Return True only for scan-batched layer dicts.
+
+    `nnx.scan` stores stacked transformer blocks under a single ``layers`` key
+    whose children are named module fields like ``attn`` and ``mlp``. In
+    contrast, ``nnx.Sequential`` also uses a ``layers`` key, but its children
+    are numeric submodule indices. Those sequential containers must not be
+    expanded.
+    """
+    if not d or all(_is_numeric_key(key) for key in d):
+        return False
+    return _infer_num_layers(d) is not None
+
+
+def expand_scanned_layers(state_dict: dict) -> dict:
+    """Expand scan-batched layer parameters into per-layer entries for HF saving.
+
+    With nnx.scan/vmap, all transformer layer parameters are stored as a single
+    batched module under a "layers" key with a leading layer dimension. This
+    expands them into separate "layers_N" sub-dicts compatible with the HF
+    format conversion utilities.
+
+    Args:
+        state_dict (dict): Model state dict potentially containing a "layers" key
+            with batched parameters.
+
+    Returns:
+        dict: State dict with "layers" expanded into "layers_0", ..., "layers_(N-1)".
+    """
+    result = {}
+    for key, value in state_dict.items():
+        if key == "layers" and isinstance(value, dict) and _should_expand_layers_dict(value):
+            num_layers = _infer_num_layers(value)
+            if num_layers is not None:
+                for i in range(num_layers):
+                    result[f"layers_{i}"] = _slice_layer(value, i)
+                continue
+        result[key] = expand_scanned_layers(value) if isinstance(value, dict) else value
+    return result

jimm/common/sharding.py

@@ -0,0 +1,98 @@
+import dataclasses
+from typing import Any, Protocol
+
+import jax
+from jax.sharding import NamedSharding
+from jax.sharding import PartitionSpec as P
+
+
+class ShardingSpec(Protocol):
+    """Protocol defining the sharding specification for model parameters.
+
+    Specs represent per-layer (non-stacked) shapes. The Transformer stacks
+    layers via nnx.vmap and patches the stacked Variable metadata to prepend
+    None for the scan axis, so the optimizer sees the correct 4-D spec.
+
+        attn_qkv_kernel  (in_features, num_heads, head_dim)
+        attn_qkv_bias    (num_heads, head_dim)
+        attn_out_kernel  (num_heads, head_dim, out_features)
+        attn_out_bias    (out_features,)
+        mlp_up_kernel    (in_features, intermediate_size)
+        mlp_up_bias      (intermediate_size,)
+        mlp_down_kernel  (intermediate_size, out_features)
+        mlp_down_bias    (out_features,)
+        layernorm        (hidden_size,)
+    """
+
+    attn_qkv_kernel: tuple[str | None, str | None, str | None]
+    attn_qkv_bias: tuple[str | None, str | None]
+    attn_out_kernel: tuple[str | None, str | None, str | None]
+    attn_out_bias: tuple[str | None]
+    mlp_up_kernel: tuple[str | None, str | None]
+    mlp_up_bias: tuple[str | None]
+    mlp_down_kernel: tuple[str | None, str | None]
+    mlp_down_bias: tuple[str | None]
+    layernorm: tuple[str | None]
+    patch_conv_kernel: tuple[str | None, str | None, str | None, str | None]
+    patch_conv_bias: tuple[str | None]
+    embed: tuple[str | None, str | None]
+    pos_embed_3d: tuple[str | None, str | None, str | None]
+    pos_embed_2d: tuple[str | None, str | None]
+    vision_pos_id: tuple[str | None, str | None]
+    text_pos_embed: tuple[str | None, str | None]
+    cls_token: tuple[str | None, str | None, str | None]
+    probe_token: tuple[str | None, str | None, str | None]
+    proj_kernel: tuple[str | None, str | None]
+    proj_bias: tuple[str | None]
+
+
+@dataclasses.dataclass(frozen=True)
+class NoSharding:
+    """No sharding - all parameters replicated."""
+
+    attn_qkv_kernel: tuple[str | None, str | None, str | None] = (None, None, None)
+    attn_qkv_bias: tuple[str | None, str | None] = (None, None)
+    attn_out_kernel: tuple[str | None, str | None, str | None] = (None, None, None)
+    attn_out_bias: tuple[str | None] = (None,)
+    mlp_up_kernel: tuple[str | None, str | None] = (None, None)
+    mlp_up_bias: tuple[str | None] = (None,)
+    mlp_down_kernel: tuple[str | None, str | None] = (None, None)
+    mlp_down_bias: tuple[str | None] = (None,)
+    layernorm: tuple[str | None] = (None,)
+    patch_conv_kernel: tuple[str | None, str | None, str | None, str | None] = (None, None, None, None)
+    patch_conv_bias: tuple[str | None] = (None,)
+    embed: tuple[str | None, str | None] = (None, None)
+    pos_embed_3d: tuple[str | None, str | None, str | None] = (None, None, None)
+    pos_embed_2d: tuple[str | None, str | None] = (None, None)
+    vision_pos_id: tuple[str | None, str | None] = (None, None)
+    text_pos_embed: tuple[str | None, str | None] = (None, None)
+    cls_token: tuple[str | None, str | None, str | None] = (None, None, None)
+    probe_token: tuple[str | None, str | None, str | None] = (None, None, None)
+    proj_kernel: tuple[str | None, str | None] = (None, None)
+    proj_bias: tuple[str | None] = (None,)
+
+
+def sharding_of(value: Any) -> NamedSharding:
+    """Returns the traced NamedSharding for a value in explicit mode."""
+
+    return jax.typeof(value).sharding
+
+
+def named_sharding_like(reference: Any, spec: P | tuple[str | None, ...]) -> NamedSharding:
+    """Builds a NamedSharding on the same mesh as the reference value."""
+
+    if not isinstance(spec, P):
+        spec = P(*spec)
+    return NamedSharding(sharding_of(reference).mesh, spec)
+
+
+def replicated_sharding(reference: Any) -> NamedSharding:
+    """Builds a replicated sharding matching the reference mesh and rank."""
+
+    return named_sharding_like(reference, P(*([None] * reference.ndim)))
+
+
+def reshard_like(value: Any, reference: Any) -> Any:
+    """Reshards a value to match the sharding of a reference value."""
+
+    return jax.sharding.reshard(value, sharding_of(reference))