# examples/vision_mamba_export.py
"""Export genuine Vision Mamba predictions through pinned official definitions.

Repository-scoped usage (CPU only; see requirements-vision-mamba.txt):
    python examples/vision_mamba_export.py --cache-dir /tmp/vision-mamba

Downloads stay in that cache. Source definitions are parsed unchanged from the
Apache-2.0 hustvl/Vim release. Credits: Lianghui Zhu et al.; models_mamba.py
also credits Facebook, Inc. (2015-present), and Mamba numerical code credits
Tri Dao and Albert Gu (2023), RMS normalization Tri Dao (2024).

Only numerical backend symbols are substituted. Official Mamba.forward keeps
use_fast_path=True and its two independently parameterized v2 directions.
The factory, model classes, residual flow and readout are unchanged. No CUDA
kernels execute. This is a reproducibility reference, not a speed benchmark.
"""

import argparse
import ast
from datetime import datetime, timezone
from functools import partial
import hashlib
import importlib.metadata
import json
import math
from pathlib import Path
import platform
import random
from types import SimpleNamespace
from typing import Optional
import urllib.request

from einops import rearrange, repeat
from PIL import Image
import torch
from torch import Tensor, nn
from torch.nn import functional as F
from timm.layers import DropPath, lecun_normal_, to_2tuple, trunc_normal_
from timm.models.vision_transformer import _cfg
from torchvision import transforms
from torchvision.models import ResNet50_Weights
from torchvision.transforms import functional as TF

import vision_mamba as teaching

REVISION = "dd0358ad1e42701f22afbefa0717cc8825cf9f45"
BASE_URL = f"https://raw.githubusercontent.com/hustvl/Vim/{REVISION}/"
HF_REVISION = "07c00e0e4ea2973d8e343afdd807128a57bc9fd5"
CHECKPOINT_URL = (
    "https://huggingface.co/hustvl/Vim-tiny-midclstok/resolve/"
    f"{HF_REVISION}/vim_t_midclstok_76p1acc.pth"
)
CHECKPOINT_SHA256 = (
    "fedb1e60267f3bc7b2a8443f2ddead8d39e019d10247a87252162b3789801fd5"
)
CHECKPOINT_BYTES = 115013118
FACTORY = "vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2"
MODEL_ID = "vim-tiny-midclstok"
MAMBA_PATH = "mamba-1p1p1/mamba_ssm/modules/mamba_simple.py"
SCAN_PATH = "mamba-1p1p1/mamba_ssm/ops/selective_scan_interface.py"
NORM_PATH = "mamba-1p1p1/mamba_ssm/ops/triton/layer_norm.py"
SOURCE_HASHES = {
    "mamba-1p1p1/benchmarks/benchmark_generation_mamba_simple.py": "fa10998acb4c53384ea71773363d351d7f90c8705b5ae696a7efa8eda678f4f9",
    "mamba-1p1p1/mamba_ssm/modules/mamba_simple.py": "2a0e9eddb5f4078e77c80ad62c40e0d2c208c12d7aa99b69cdcbabbe5b1673a0",
    "mamba-1p1p1/mamba_ssm/ops/triton/layer_norm.py": "b7a77c2f305650382b30509d503a3536aa09fd28b6b4595d7c2160813a3f4317",
    "vim/rope.py": "b05a75d5b2246c0b758f8f941f20c1fc0943d77fa778dd9f79d00a2b4e4edd91",
    "vim/scripts/ft-vim-t.sh": "dbeca4aa8e7cf527999db1e79c290557b8e37f5c018a349e11196b2a89db2af0",
    "README.md": "b7fd237a4783e8e1eaa19c4e2de864e2bfdba152685ced9e48ec0564ace58b4a",
    "LICENSE.txt": "760939b000194d04548ede6a857bbe735d1695d8422ec85955c8e2bd7f4b95c5",
    "vim/models_mamba.py": "f32a57266f9c5e2824f51cded579f6a5eb722001d5f1cad3ea7e7179ab55f0f5",
    "vim/datasets.py": "86a71891ee80ce916071c712ca499fa2cab409aeb60283ff1c8bb3cc761beb86",
    "vim/main.py": "4c30749ac23bd3f827f1146c04ac8f98657ee23f1cb0b2ddde1cbec6bba8bc90",
    "vim/scripts/pt-vim-t.sh": "0e7c5231dbb97df402246aab6334adac94e92d94f263f868b5222e6e5494e5b5",
    "mamba-1p1p1/mamba_ssm/ops/selective_scan_interface.py": "c0bb4ffb85fdbd1a0700ba27198638c25c929ad83a4ddd6f50c69357455b093f",
}
TRANSFORM = {
    "resizeSize": 256,
    "cropSize": 224,
    "evalCropRatio": 0.875,
    "interpolation": "bicubic",
    "antialias": True,
    "mean": [0.485, 0.456, 0.406],
    "std": [0.229, 0.224, 0.225],
    "resizeRule": "shorter edge; preserve aspect ratio",
    "cropRule": "torchvision CenterCrop; round((dimension - 224) / 2)",
    "colorMode": "RGB",
    "tensorScale": "uint8 / 255",
    "source": BASE_URL + "vim/datasets.py",
    "settingsSource": BASE_URL + "vim/main.py",
}


def digest(path):
    """Hash exact file bytes."""
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def verify_file(path, expected):
    """Reject cached inputs whose bytes differ from the reviewed revision."""
    if digest(path) != expected:
        raise ValueError(f"SHA-256 mismatch: {path}")


def cached_download(url, path, expected):
    """Download and verify before making an artifact available atomically."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if not path.exists():
        temporary = path.with_suffix(path.suffix + ".download")
        with urllib.request.urlopen(
            url, timeout=60
        ) as response, temporary.open("wb") as output:
            while chunk := response.read(1024 * 1024):
                output.write(chunk)
        verify_file(temporary, expected)
        temporary.replace(path)
    verify_file(path, expected)
    return path


def execute_definitions(path, names, namespace):
    """Compile selected complete AST definitions without editing their bodies.

    Avoiding import-time CUDA dependencies does not rewrite the selected code.
    Decorators are preserved; register_model only skips global timm registration.
    """
    tree = ast.parse(Path(path).read_text())
    selected = [
        node
        for node in tree.body
        if isinstance(node, (ast.FunctionDef, ast.ClassDef))
        and node.name in names
    ]
    if {node.name for node in selected} != set(names):
        raise ValueError(f"Missing official definitions in {path}")
    exec(
        compile(ast.Module(body=selected, type_ignores=[]), str(path), "exec"),
        namespace,
    )


def cpu_causal_conv1d(x, weight, bias=None, activation=None):
    """Replace CUDA causal-conv1d with exact grouped PyTorch left padding."""
    result = F.conv1d(
        F.pad(x, (weight.shape[-1] - 1, 0)),
        weight[:, None],
        bias,
        groups=x.shape[1],
    )
    if activation not in (None, "silu", "swish"):
        raise ValueError("Unsupported causal convolution activation")
    return result if activation is None else F.silu(result)


def load_official_definitions(upstream):
    """Load unchanged classifier/factory and numerical reference definitions."""
    upstream = Path(upstream)
    for path, expected in SOURCE_HASHES.items():
        verify_file(upstream / path, expected)
    namespace = {
        "__name__": "vision_mamba_official_pinned",
        "torch": torch,
        "nn": nn,
        "F": F,
        "Tensor": Tensor,
        "Optional": Optional,
        "math": math,
        "partial": partial,
        "random": random,
        "rearrange": rearrange,
        "repeat": repeat,
        "DropPath": DropPath,
        "to_2tuple": to_2tuple,
        "trunc_normal_": trunc_normal_,
        "lecun_normal_": lecun_normal_,
        "register_model": lambda function: function,
        "_cfg": _cfg,
        "causal_conv1d_fn": cpu_causal_conv1d,
    }
    execute_definitions(
        upstream / SCAN_PATH,
        ["selective_scan_ref", "mamba_inner_ref"],
        namespace,
    )
    namespace["selective_scan_fn"] = namespace["selective_scan_ref"]
    execute_definitions(
        upstream / NORM_PATH, ["rms_norm_ref", "RMSNorm"], namespace
    )

    def rms_norm_cpu(
        x,
        weight,
        bias,
        residual=None,
        eps=1e-5,
        prenorm=False,
        residual_in_fp32=False,
        dropout_p=0.0,
    ):
        # The selected model executes float32 only. Upcast preserves carried
        # residual semantics and calls the untouched official PyTorch reference.
        return namespace["rms_norm_ref"](
            x,
            weight,
            bias,
            residual=residual,
            eps=eps,
            prenorm=prenorm,
            upcast=residual_in_fp32,
            dropout_p=dropout_p,
        )

    def inner_without_projection(
        xz,
        conv_weight,
        conv_bias,
        x_weight,
        dt_weight,
        a,
        b=None,
        c=None,
        d=None,
        delta_bias=None,
        delta_softplus=True,
    ):
        # Upstream offers mamba_inner_ref with an output projection but no
        # no-output-projection reference. An identity projection and transpose
        # expose the exact pre-projection tensor without changing its body.
        identity = torch.eye(a.shape[0], dtype=xz.dtype, device=xz.device)
        result = namespace["mamba_inner_ref"](
            xz,
            conv_weight,
            conv_bias,
            x_weight,
            dt_weight,
            identity,
            None,
            a,
            b,
            c,
            d,
            delta_bias=delta_bias,
            delta_softplus=delta_softplus,
        )
        return result.transpose(1, 2)

    namespace["rms_norm_fn"] = rms_norm_cpu
    namespace["layer_norm_fn"] = None  # Never selected by this RMSNorm factory.
    namespace["mamba_inner_fn_no_out_proj"] = inner_without_projection
    execute_definitions(upstream / MAMBA_PATH, ["Mamba"], namespace)
    execute_definitions(
        upstream / "vim/models_mamba.py",
        [
            "PatchEmbed",
            "Block",
            "create_block",
            "_init_weights",
            "segm_init_weights",
            "VisionMamba",
            FACTORY,
        ],
        namespace,
    )
    return namespace


def load_checkpoint(official, path):
    """Strictly load the recorded model key, allowing only Namespace metadata."""
    verify_file(path, CHECKPOINT_SHA256)
    if path.stat().st_size != CHECKPOINT_BYTES:
        raise ValueError("Checkpoint byte count mismatch")
    with torch.serialization.safe_globals([argparse.Namespace]):
        checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model = official[FACTORY](pretrained=False, drop_path_rate=0.0)
    result = model.load_state_dict(checkpoint["model"], strict=True)
    model.eval()
    if not all(
        layer.mixer.use_fast_path and layer.mixer.bimamba_type == "v2"
        for layer in model.layers
    ):
        raise ValueError(
            "Official bidirectional v2 fast-path control flow required"
        )
    return model, {
        "strict": True,
        "missingKeys": list(result.missing_keys),
        "unexpectedKeys": list(result.unexpected_keys),
        "checkpointKeys": list(checkpoint),
        "stateDictKeys": len(checkpoint["model"]),
        "checkpointEpoch": checkpoint["epoch"],
    }


def compare_teaching_core(official):
    """Compare outputs and all input/parameter gradients with official code."""
    torch.manual_seed(7)
    args = [
        torch.randn(2, 8, 9),
        torch.randn(2, 8, 9),
        -torch.rand(8, 3),
        torch.randn(2, 3, 9),
        torch.randn(2, 3, 9),
        torch.randn(8),
        torch.randn(2, 8, 9),
        torch.randn(8),
    ]
    ours = [value.clone().requires_grad_() for value in args]
    theirs = [value.clone().requires_grad_() for value in args]
    actual = teaching.selective_scan(*ours, delta_softplus=True)
    expected = official["selective_scan_ref"](*theirs, delta_softplus=True)
    torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5)
    actual.square().mean().backward()
    expected.square().mean().backward()
    errors = []
    for left, right in zip(ours, theirs):
        torch.testing.assert_close(left.grad, right.grad, atol=2e-4, rtol=2e-4)
        errors.append(float((left.grad - right.grad).abs().max()))
    scan_error = float((actual - expected).abs().max().detach())
    # Independent entire mixer comparison checks reverse conv/order/gate and
    # both sets of directional parameter gradients, beyond scan-only agreement.
    reference = official["Mamba"](
        8, d_state=3, bimamba_type="v2", if_divide_out=True
    )
    model = teaching.BidirectionalMixer(8, state_size=3)
    pairs = [
        (model.in_proj.weight, reference.in_proj.weight),
        (model.out_proj.weight, reference.out_proj.weight),
    ]
    for direction, suffix in [
        (model.forward_direction, ""),
        (model.backward_direction, "_b"),
    ]:
        pairs.extend(
            [
                (
                    direction.conv.weight,
                    getattr(reference, "conv1d" + suffix).weight,
                ),
                (
                    direction.conv.bias,
                    getattr(reference, "conv1d" + suffix).bias,
                ),
                (
                    direction.x_proj.weight,
                    getattr(reference, "x_proj" + suffix).weight,
                ),
                (
                    direction.dt_proj.weight,
                    getattr(reference, "dt_proj" + suffix).weight,
                ),
                (
                    direction.dt_proj.bias,
                    getattr(reference, "dt_proj" + suffix).bias,
                ),
                (direction.a_log, getattr(reference, "A" + suffix + "_log")),
                (direction.d, getattr(reference, "D" + suffix)),
            ]
        )
    with torch.no_grad():
        for left, right in pairs:
            left.copy_(right)
    x = torch.randn(2, 9, 8, requires_grad=True)
    other_x = x.detach().clone().requires_grad_()
    actual = model(x)
    expected = reference(other_x)
    torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5)
    mixer_error = float((actual - expected).abs().max().detach())
    actual.square().mean().backward()
    expected.square().mean().backward()
    for left, right in [(x, other_x), *pairs]:
        torch.testing.assert_close(left.grad, right.grad, atol=2e-4, rtol=2e-4)
        errors.append(float((left.grad - right.grad).abs().max()))
    return {
        "passed": True,
        "seed": 7,
        "reference": "unchanged selective_scan_ref and official Mamba.forward v2 with mamba_inner_ref CPU backend",
        "classifier": "unchanged official VisionMamba class and factory executed directly; no authored pretrained port",
        "scanInputShape": [2, 8, 9],
        "mixerInputShape": [2, 9, 8],
        "scanMaxAbsoluteError": scan_error,
        "mixerMaxAbsoluteError": mixer_error,
        "maxOutputAbsoluteError": max(scan_error, mixer_error),
        "maxGradientAbsoluteError": max(errors),
        "outputAtol": 2e-5,
        "outputRtol": 2e-5,
        "gradientAtol": 2e-4,
        "gradientRtol": 2e-4,
        "gradientCoverage": "all 8 scan inputs; mixer input and all 16 parameter tensors",
    }


def preprocess(image):
    """Return the exact crop, normalized tensor and integer crop geometry."""
    resized = TF.resize(
        image.convert("RGB"),
        256,
        interpolation=transforms.InterpolationMode.BICUBIC,
        antialias=True,
    )
    crop = TF.center_crop(resized, [224, 224])
    left = round((resized.width - 224) / 2)
    top = round((resized.height - 224) / 2)
    tensor = TF.normalize(
        TF.to_tensor(crop), TRANSFORM["mean"], TRANSFORM["std"]
    )
    return (
        crop,
        tensor.unsqueeze(0),
        {
            "resizedShape": [resized.height, resized.width],
            "cropBox": [left, top, left + 224, top + 224],
        },
    )


def official_transform(path):
    """Execute unchanged pinned build_transform(False) with release eval args."""
    namespace = {
        "transforms": transforms,
        "IMAGENET_DEFAULT_MEAN": (0.485, 0.456, 0.406),
        "IMAGENET_DEFAULT_STD": (0.229, 0.224, 0.225),
    }
    execute_definitions(path, ["build_transform"], namespace)
    return namespace["build_transform"](
        False, SimpleNamespace(input_size=224, eval_crop_ratio=0.875)
    )


def summarize_logits(logits, categories):
    """Store all logits/probabilities, sorting ties by ascending class index."""
    if (
        logits.shape != (1000,)
        or len(categories) != 1000
        or not torch.isfinite(logits).all()
    ):
        raise ValueError("Expected exactly 1000 finite logits and labels")
    probabilities = logits.softmax(dim=0)
    indices = sorted(
        range(1000), key=lambda index: (-float(logits[index]), index)
    )[:5]
    return {
        "logits": logits.tolist(),
        "probabilities": probabilities.tolist(),
        "top5": [
            {
                "index": index,
                "label": categories[index],
                "logit": float(logits[index]),
                "probability": float(probabilities[index]),
            }
            for index in indices
        ],
    }


def record_forward(model, tensor):
    """Capture actual tensors: mixer contribution differs from residual sum."""
    stages = []
    handles = []

    def record(name, value, semantics, input_shape=None, **extra):
        stages.append(
            {
                "name": name,
                "outputShape": list(value.shape),
                "inputShape": input_shape or list(value.shape),
                "semantics": semantics,
                **extra,
            }
        )

    def patches_hook(_module, inputs, output):
        record(
            "patches",
            output,
            "Row-major projected patch tokens before CLS and position addition",
            list(inputs[0].shape),
        )

    def tokens_hook(_module, inputs):
        record(
            "tokens",
            inputs[0],
            "Middle CLS inserted at 98; learned absolute positions added",
        )

    def block_hook(index):
        def hook(_module, inputs, output):
            contribution, carried = output
            record(
                f"block{index + 1}-mixer",
                contribution,
                "Mixer contribution only; not the full residual feature",
                list(inputs[0].shape),
                block=index + 1,
            )
            record(
                f"block{index + 1}-residual",
                carried + contribution,
                "Full residual feature after adding this block's mixer contribution",
                block=index + 1,
            )

        return hook

    def head_hook(_module, inputs, output):
        record(
            "normalized-cls",
            inputs[0],
            "Final RMSNorm applied after last residual addition, then middle CLS selected",
            list(inputs[0].shape),
        )
        record(
            "logits",
            output,
            "ImageNet-1K linear classifier logits",
            list(inputs[0].shape),
        )

    try:
        handles.append(model.patch_embed.register_forward_hook(patches_hook))
        handles.append(model.layers[0].register_forward_pre_hook(tokens_hook))
        for index in [0, 11, 23]:
            handles.append(
                model.layers[index].register_forward_hook(block_hook(index))
            )
        handles.append(model.head.register_forward_hook(head_hook))
        with torch.inference_mode():
            logits = model(tensor)[0]
    finally:
        for handle in handles:
            handle.remove()
    return logits, stages


def export(cache, root, threads=2):
    """Write predictions only after all measured examples and oracles succeed."""
    cache.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(threads)
    torch.manual_seed(0)
    torch.use_deterministic_algorithms(True)
    upstream = cache / "upstream"
    artifacts = []
    for path, sha in SOURCE_HASHES.items():
        downloaded = cached_download(BASE_URL + path, upstream / path, sha)
        artifacts.append(
            {"path": path, "url": BASE_URL + path, "sha256": digest(downloaded)}
        )
    checkpoint = cached_download(
        CHECKPOINT_URL, cache / "vim_t_midclstok_76p1acc.pth", CHECKPOINT_SHA256
    )
    official = load_official_definitions(upstream)
    oracle = compare_teaching_core(official)
    model, strict_load = load_checkpoint(official, checkpoint)
    parameter_count = sum(parameter.numel() for parameter in model.parameters())
    print(
        json.dumps(
            {
                "parameterCount": parameter_count,
                "strictLoad": strict_load,
                "oracle": oracle,
            }
        ),
        flush=True,
    )
    transform = official_transform(upstream / "vim/datasets.py")
    photos = json.loads(
        (root / "examples/efficientnet-sample-images/manifest.json").read_text()
    )["photos"]
    categories = ResNet50_Weights.IMAGENET1K_V2.meta["categories"]
    output_dir = root / "public/images/papers/vision-mamba/examples"
    output_dir.mkdir(parents=True, exist_ok=True)
    runs = []
    for photo in photos:
        original = root / "public" / photo["originalImage"].lstrip("/")
        verify_file(original, photo["source"]["sha256"])
        with Image.open(original) as image:
            crop, tensor, geometry = preprocess(image)
            torch.testing.assert_close(
                tensor[0], transform(image.convert("RGB")), atol=0, rtol=0
            )
        temporary = output_dir / f".{photo['id']}-crop.png"
        crop.save(temporary, format="PNG")
        crop_hash = digest(temporary)
        crop_path = output_dir / f"{photo['id']}-{crop_hash[:12]}.png"
        temporary.replace(crop_path)
        print(
            f"Executing official Vim-T on {photo['id']} (CPU float32)",
            flush=True,
        )
        logits, stages = record_forward(model, tensor)
        run = {
            "id": photo["id"] + "-" + MODEL_ID,
            "photoId": photo["id"],
            "modelId": MODEL_ID,
            "inputImage": "/" + str(crop_path.relative_to(root / "public")),
            "inputSha256": crop_hash,
            "inputShape": list(tensor.shape),
            "normalizedTensorSha256": hashlib.sha256(
                tensor.numpy().tobytes()
            ).hexdigest(),
            "preprocessingReferenceMaxAbsoluteError": 0.0,
            **geometry,
            **summarize_logits(logits, categories),
            "stages": stages,
        }
        runs.append(run)
        print(
            json.dumps({"photo": photo["id"], "top5": run["top5"]}), flush=True
        )
    manifest = {
        "schemaVersion": 1,
        "provenance": {
            "sourceRevision": REVISION,
            "sourceUrl": BASE_URL + "vim/models_mamba.py",
            "referenceSha256": SOURCE_HASHES["vim/models_mamba.py"],
            "exportSha256": digest(Path(__file__)),
            "teachingSha256": digest(root / "examples/vision_mamba.py"),
            "exportedAt": datetime.now(timezone.utc).isoformat(),
            "device": "cpu",
            "dtype": "float32",
            "numThreads": threads,
            "deterministicAlgorithms": True,
            "backend": "unchanged official classes with explicit PyTorch CPU numerical reference backends",
            "backendSubstitutions": [
                "mamba_inner_fn_no_out_proj -> unchanged mamba_inner_ref with identity out projection then BLD-to-BDL transpose",
                "selective_scan_fn -> unchanged official selective_scan_ref",
                "causal_conv1d_fn -> torch.nn.functional.conv1d with left padding, groups=channels, then SiLU",
                "rms_norm_fn -> unchanged official rms_norm_ref; residual_in_fp32 maps to upcast=True",
                "AST extraction avoids CUDA/Triton imports; original class/factory/function bodies and decorators are preserved",
                "register_model is an identity decorator to avoid global timm registry mutation; real timm initialization, DropPath and _cfg are supplied",
            ],
            "versions": {
                "python": platform.python_version(),
                **{
                    name: importlib.metadata.version(name)
                    for name in [
                        "torch",
                        "torchvision",
                        "Pillow",
                        "timm",
                        "numpy",
                        "einops",
                    ]
                },
            },
            "platform": platform.platform(),
            "oracleAgreement": oracle,
            "preprocessingReference": "unchanged pinned build_transform(False), input_size=224, eval_crop_ratio=0.875; bitwise tensor equality on every photo",
            "sourceArtifacts": artifacts,
            "codeLicense": {
                "name": "Apache-2.0",
                "url": BASE_URL + "LICENSE.txt",
            },
            "labelsSource": "torchvision.models.ResNet50_Weights.IMAGENET1K_V2.meta.categories (vocabulary only)",
        },
        "categories": categories,
        "models": [
            {
                "id": MODEL_ID,
                "label": "Vim-T (middle CLS, baseline)",
                "factory": FACTORY,
                "strictLoad": strict_load,
                "parameterCount": parameter_count,
                "weights": {
                    "url": CHECKPOINT_URL,
                    "revision": HF_REVISION,
                    "sha256": digest(checkpoint),
                    "bytes": checkpoint.stat().st_size,
                    "key": "model",
                },
                "config": {
                    "inputSize": 224,
                    "patchSize": 16,
                    "stride": 16,
                    "patchGrid": [14, 14],
                    "patchTokens": 196,
                    "classTokenIndex": 98,
                    "sequenceLength": 197,
                    "embedDim": 192,
                    "depth": 24,
                    "innerDim": 384,
                    "stateSize": 16,
                    "dtRank": 12,
                    "convKernel": 4,
                    "bimambaType": "v2",
                    "divideOut": True,
                    "useFastPath": True,
                    "rmsNorm": True,
                    "normEpsilon": 1e-5,
                    "residualInFp32": True,
                    "fusedAddNorm": True,
                    "dropPathRate": 0.0,
                    "absolutePositions": True,
                    "rope": False,
                    "classToken": "middle",
                    "finalPoolType": "mean",
                    "effectiveReadout": "normalized middle class token (takes precedence over finalPoolType)",
                    "numClasses": 1000,
                },
                "transform": TRANSFORM,
            }
        ],
        "photos": photos,
        "runs": runs,
        "limitations": [
            "Three individual photographs, not an accuracy evaluation.",
            "No latency, calibration, attention-map or saliency measurement.",
            "CPU float32 reference results; released GPU fused kernels may differ numerically.",
            "Recorded stage shapes come from real tensors; schematic planes do not visualize activation values.",
            "The teaching model is random-weight and separate from this official pretrained classifier.",
        ],
    }
    target = root / "public/data/vision-mamba/predictions.json"
    target.parent.mkdir(parents=True, exist_ok=True)
    temporary = target.with_suffix(".json.tmp")
    temporary.write_text(json.dumps(manifest, indent=2, allow_nan=False) + "\n")
    temporary.replace(target)
    print(f"Wrote {target}", flush=True)


def main():
    """Parse reproducibility controls; weights and architecture remain pinned."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache-dir", type=Path, required=True)
    parser.add_argument(
        "--repo-root", type=Path, default=Path(__file__).resolve().parents[1]
    )
    parser.add_argument("--threads", type=int, default=2)
    args = parser.parse_args()
    if args.threads < 1:
        parser.error("--threads must be positive")
    export(args.cache_dir, args.repo_root, args.threads)


if __name__ == "__main__":
    main()