export/export.py

# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""Export EmbeddingGemma 2 (text + images + audio) to the three Core AI `.aimodel` assets of this package.

    # 1. download the original checkpoint: hf download google/embeddinggemma-2 --local-dir ./embeddinggemma-2
    # 2. python3 -W ignore export.py --model-dir ./embeddinggemma-2 --out <bundle folder> \
    #        --buckets 32,64,128,288,512,1024,2048 --qchunks 12 --audio-buckets 192,384,1024,1536,3000 --date 2026-10-12
    #    (the released bundle, ~3 min)

Written to <out> (see README "Why three assets"):
  embeddinggemma-2-audio.aimodel     functions audio_<Tb>: features [1,Tb,128] fp16 (log-mel frames), frame_mask [1,Tb],
                                     half_mask [1,Tb/2], token_mask [1,Tb/4] fp16 -> soft_tokens [1,Tb/4,512] fp16
                                     (eg2_audio.py; Tb in 192, 384, 1024, 1536, 3000 frames of 10 ms)
  embeddinggemma-2-text-mm.aimodel   functions mm_<S>: input_ids [1,S] int32, attention_mask [1,S] fp16,
                                     mm_embeds [1,S,512] fp16, mm_mask [1,S] fp16 -> embedding [1,768] fp16.
                                     The text encoder (eg2_text.py, unchanged from v1.0) with a soft-token mix at the
                                     input (eg2_text_mm.py); with zero soft tokens it is exactly the text encoder.
  embeddinggemma-2-vision.aimodel    function vision_2520: patches [1,2520,768] fp16 (pixel values 0..255), pos_x/pos_y [1,2520] int32,
                                     cos/sin [1,1,2520,64] fp16, patch_mask [1,2520] fp16 -> soft_tokens [1,280,512] fp16
                                     (eg2_vision.py)
  tokenizer/                         unchanged Gemma tokenizer (text queries and the image sequence ids)
  config.json, MODEL-ATTRIBUTION.json

`--part text` / `--part vision` / `--part audio` re-exports one asset and keeps the others' entries of an existing
config.json.

The result is functionally identical to the released bundle, not byte-identical: the converter names internal graph
resources at random (main.mlirb / main.hash differ) and the runtime writes metadata.json with the export time.
config.json, MODEL-ATTRIBUTION.json (with --date 2026-10-12) and the tokenizer files come out byte-identical.

Requirements: Python 3.12, torch 2.9.0, coreai-torch 0.4.3, coreai-opt 0.3.0, torchao 0.17.0, coreai-core 1.0.0b3,
safetensors (see ../python/requirements.lock).
"""
from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import subprocess
import sys
import time
from pathlib import Path

import torch

LAB = Path(__file__).resolve().parent
sys.path.insert(0, str(LAB))
import eg2_audio  # noqa: E402
import eg2_text  # noqa: E402
import eg2_text_mm  # noqa: E402
import eg2_vision  # noqa: E402

T0 = time.time()
PROMPTS = {
    "document": "title: none | text: ",
    "search_query": "task: search result | query: ",
    "question_answering": "task: question answering | query: ",
    "fact_checking": "task: fact checking | query: ",
    "code_retrieval": "task: code retrieval | query: ",
    "classification": "task: classification | query: ",
    "clustering": "task: clustering | query: ",
    "sentence_similarity": "task: sentence similarity | query: ",
}
SPECIAL = {"bos": 2, "eos": 1, "pad": 0, "boi": 255999, "image": 258880, "eoi": 258882, "video": 258884,
           "boa": 256000, "audio": 258881, "eoa": 258883}


def log(msg: str) -> None:
    print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True)


def sha256_file(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 24), b""):
            h.update(chunk)
    return h.hexdigest()


def dir_bytes(p: Path) -> int:
    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) if p.is_dir() else p.stat().st_size


def save_asset(prog, path: Path, description: str) -> None:
    from coreai.runtime import AIModelAssetMetadata

    shutil.rmtree(path, ignore_errors=True)
    meta = AIModelAssetMetadata()
    meta.author = "Gemma Team (Google); Core AI port"
    meta.license = "Apache-2.0"
    meta.model_description = description
    meta.creation_date = int(time.time())
    prog.save_asset(path, meta)
    log(f"saved {path} ({dir_bytes(path) / 2**20:.1f} MiB)")
    try:
        print(subprocess.run(["xcrun", "coreai-build", "inspect", str(path), "--no-metadata"], capture_output=True, text=True, timeout=120).stdout[:2500])
    except Exception as e:  # noqa: BLE001
        log(f"inspect failed: {e}")


def export_vision(args, out: Path) -> dict:
    import coreai_torch

    eg2_vision.PRECISION_MODE = "ane"
    eg2_vision.QUERY_CHUNKS = args.qchunks
    model = eg2_vision.load_model(args.model_dir, torch.float16)
    if args.num_layers:
        model.layers = model.layers[: args.num_layers]
    log(f"vision model loaded: {sum(p.numel() for p in model.parameters()) / 1e6:.1f} M params, qchunks={args.qchunks}")
    N = eg2_vision.MAX_PATCHES
    ex = {"patches": torch.rand(1, N, 768, dtype=torch.float16), "pos_x": torch.zeros(1, N, dtype=torch.int32),
          "pos_y": torch.zeros(1, N, dtype=torch.int32), "cos": torch.ones(1, 1, N, 64, dtype=torch.float16),
          "sin": torch.zeros(1, 1, N, 64, dtype=torch.float16), "patch_mask": torch.ones(1, N, dtype=torch.float16)}
    with torch.no_grad():
        ep = torch.export.export(model, args=(), kwargs=ex)
    log("torch.export (vision) done")
    ep = ep.run_decompositions(coreai_torch.get_decomp_table())
    conv = coreai_torch.TorchConverter(mode=coreai_torch.TorchConverter.Mode.DEBUG if args.debug_info else coreai_torch.TorchConverter.Mode.RELEASE)
    fn = f"vision_{N}"
    conv.add_exported_program(ep, input_names=["patches", "pos_x", "pos_y", "cos", "sin", "patch_mask"], output_names=["soft_tokens"], entrypoint_name=fn)
    prog = conv.to_coreai()
    log("converted vision to AIProgram")
    asset = out / "embeddinggemma-2-vision.aimodel"
    save_asset(prog, asset, f"EmbeddingGemma 2 vision encoder ({args.hf_id}): 2520 patches -> 280 soft tokens in the text space; "
                            f"fp16, integer pixel inputs (0..255), cell-major pooling, max-abs RMSNorm.")
    return {"asset": asset.name, "function": fn, "max_patches": N, "max_soft_tokens": eg2_vision.MAX_SOFT, "qchunks": args.qchunks}


def export_audio(args, out: Path, buckets: list[int]) -> dict:
    """One dynamic-T torch.export (T = 4t) of eg2_audio, specialized to the frame buckets as functions audio_<Tb>."""
    import coreai_torch
    from torch.export import Dim

    eg2_audio.PRECISION_MODE = "ane"
    model = eg2_audio.load_model(args.model_dir, torch.float16)
    if args.num_layers:
        model.layers = model.layers[: args.num_layers]
    log(f"audio model loaded: {sum(p.numel() for p in model.parameters()) / 1e6:.1f} M params; frame buckets {buckets}")
    assert all(b % 4 == 0 for b in buckets) and len(buckets) >= 2
    names = ["features", "frame_mask", "half_mask", "token_mask"]
    b0 = buckets[1]
    ex = {"features": torch.randn(1, b0, 128, dtype=torch.float16), "frame_mask": torch.ones(1, b0, dtype=torch.float16),
          "half_mask": torch.ones(1, b0 // 2, dtype=torch.float16), "token_mask": torch.ones(1, b0 // 4, dtype=torch.float16)}
    t = Dim("t", min=buckets[0] // 4, max=buckets[-1] // 4)
    dyn = {"features": {1: 4 * t}, "frame_mask": {1: 4 * t}, "half_mask": {1: 2 * t}, "token_mask": {1: t}}
    with torch.no_grad():
        ep = torch.export.export(model, args=(), kwargs=ex, dynamic_shapes=dyn)
    log("torch.export (audio) done")
    ep = ep.run_decompositions(coreai_torch.get_decomp_table())
    conv = coreai_torch.TorchConverter(mode=coreai_torch.TorchConverter.Mode.DEBUG if args.debug_info else coreai_torch.TorchConverter.Mode.RELEASE)
    conv.add_exported_program(ep, input_names=names, output_names=["soft_tokens"], entrypoint_name="audio")
    prog = conv.to_coreai()
    log("converted audio to AIProgram")
    cfg = {f'"{b}"': {"features": (1, b, 128), "frame_mask": (1, b), "half_mask": (1, b // 2), "token_mask": (1, b // 4)} for b in buckets}
    prog.set_static_shape_config("audio", cfg)
    functions = [f"audio_{b}" for b in buckets]
    log(f"static shape specializations: {functions}")
    asset = out / "embeddinggemma-2-audio.aimodel"
    save_asset(prog, asset, f"EmbeddingGemma 2 audio encoder ({args.hf_id}): log-mel frames -> soft tokens in the text space; "
                            f"fp16, banded local attention, max-abs RMSNorm; frame buckets {buckets}.")
    return {"asset": asset.name, "functions": functions, "frame_buckets": buckets, "max_frames": buckets[-1],
            "max_soft_tokens": buckets[-1] // 4}


def export_text_mm(args, out: Path, buckets: list[int]) -> dict:
    import coreai_torch
    from torch.export import Dim

    eg2_text.PRECISION_MODE = "ane"
    max_seq = max(buckets)
    model = eg2_text_mm.load_model(args.model_dir, torch.float16, max_seq=max_seq)
    if args.num_layers:
        model.layers = model.layers[: args.num_layers]
        model.cfg.layer_types = model.cfg.layer_types[: args.num_layers]
    log(f"text-mm model loaded: {sum(p.numel() for p in model.parameters()) / 1e6:.1f} M params")
    S0 = buckets[min(1, len(buckets) - 1)]
    ref = {"input_ids": torch.zeros(1, S0, dtype=torch.int32), "attention_mask": torch.ones(1, S0, dtype=torch.float16),
           "mm_embeds": torch.zeros(1, S0, 512, dtype=torch.float16), "mm_mask": torch.zeros(1, S0, dtype=torch.float16)}
    ref["input_ids"][0, :8] = torch.tensor([2, 8071, 236787, 3927, 1354, 1109, 7609, 1], dtype=torch.int32)
    ref["attention_mask"][0, 8:] = 0
    S = Dim("S", min=min(8, buckets[0]), max=max_seq)
    dyn = {"input_ids": {1: S}, "attention_mask": {1: S}, "mm_embeds": {1: S}, "mm_mask": {1: S}}
    with torch.no_grad():
        ep = torch.export.export(model, args=(), kwargs=ref, dynamic_shapes=dyn)
    log("torch.export (text-mm) done")
    ep = ep.run_decompositions(coreai_torch.get_decomp_table())
    conv = coreai_torch.TorchConverter(mode=coreai_torch.TorchConverter.Mode.DEBUG if args.debug_info else coreai_torch.TorchConverter.Mode.RELEASE)
    conv.add_exported_program(ep, input_names=["input_ids", "attention_mask", "mm_embeds", "mm_mask"], output_names=["embedding"], entrypoint_name="mm")
    prog = conv.to_coreai()
    log("converted text-mm to AIProgram")
    cfg = {f'"{b}"': {"input_ids": (1, b), "attention_mask": (1, b), "mm_embeds": (1, b, 512), "mm_mask": (1, b)} for b in buckets}
    prog.set_static_shape_config("mm", cfg)
    functions = [f"mm_{b}" for b in buckets]
    log(f"static shape specializations: {functions}")
    asset = out / "embeddinggemma-2-text-mm.aimodel"
    save_asset(prog, asset, f"EmbeddingGemma 2 text encoder ({args.hf_id}) with soft-token (image) inputs, mean pooling, 768-d "
                            f"projection and L2 normalization in-graph; fp16 (Neural Engine); sequence buckets {buckets}.")
    return {"asset": asset.name, "functions": functions, "buckets": buckets, "max_tokens": max_seq}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--model-dir", required=True, help="folder with the original google/embeddinggemma-2 checkpoint")
    ap.add_argument("--hf-id", default="google/embeddinggemma-2")
    ap.add_argument("--out", required=True)
    ap.add_argument("--part", default="all", choices=["all", "text", "vision", "audio"])
    ap.add_argument("--buckets", default="32,64,128,288,512,1024,2048")
    ap.add_argument("--qchunks", type=int, default=12,
                    help="split the vision attention's queries into N blocks (released: 12; 1 = one 2520x2520 matmul per layer)")
    ap.add_argument("--audio-buckets", default="192,384,1024,1536,3000",
                    help="frame buckets of the audio encoder (multiples of 4; only the default set is validated on the Neural "
                         "Engine, see README)")
    ap.add_argument("--num-layers", type=int, default=None, help="smoke export with the first N layers")
    ap.add_argument("--debug-info", action="store_true")
    ap.add_argument("--date", default=None, help="port date written to MODEL-ATTRIBUTION.json (default: today)")
    args = ap.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    buckets = sorted(int(b) for b in args.buckets.split(","))
    log(f"args {vars(args)}")
    existing = json.loads((out / "config.json").read_text()) if (out / "config.json").exists() else {}
    vision = ({k: existing["vision"][k] for k in ("asset", "function", "max_patches", "max_soft_tokens", "qchunks")}
              if "vision" in existing else None)
    text = {k: existing[k] for k in ("asset", "functions", "buckets", "max_tokens")} if "asset" in existing else None
    audio = ({k: existing["audio"][k] for k in ("asset", "functions", "frame_buckets", "max_frames", "max_soft_tokens")}
             if "audio" in existing else None)
    audio_buckets = sorted(int(b) for b in args.audio_buckets.split(","))
    if args.part in ("all", "text"):
        text = export_text_mm(args, out, buckets)
    if args.part in ("all", "vision"):
        vision = export_vision(args, out)
    if args.part in ("all", "audio"):
        audio = export_audio(args, out, audio_buckets)
    if text is None or vision is None or audio is None:
        raise SystemExit("--part text|vision|audio needs an existing config.json in --out with the other assets' entries")

    tok_dir = out / "tokenizer"
    tok_dir.mkdir(exist_ok=True)
    for f in ("tokenizer.json", "tokenizer_config.json"):
        shutil.copy(Path(args.model_dir) / f, tok_dir / f)
    cfg_json = {
        "base_model": args.hf_id, "format": "coreai", **text, "precision_mode": "ane",
        "inputs": {
            "input_ids": "[1, S] int32 (bos=2 ... eos=1, right-padded with pad=0; image slots hold 258880)",
            "attention_mask": "[1, S] float16 (1 = token, 0 = padding)",
            "mm_embeds": "[1, S, 512] float16 soft tokens at their slots (zeros elsewhere), unscaled",
            "mm_mask": "[1, S] float16 (1 at a soft-token slot)"},
        "output": {"embedding": "[1, 768] float16, L2-normalized; Matryoshka: keep the first 512/256/128 and re-normalize"},
        "text_only_use": "mm_embeds = 0 and mm_mask = 0 give the text embedding (bit-identical to the v1.0 main_<S> functions)",
        "vision": {**vision, "inputs": {
            "patches": "[1, 2520, 768] float16: 16x16x3 patches (py, px, c), integer pixel values 0..255 (not /255), CELL-MAJOR order, zero padded",
            "pos_x": "[1, 2520] int32 patch column (0 for padding)", "pos_y": "[1, 2520] int32 patch row (0 for padding)",
            "cos": "[1, 1, 2520, 64] float16 axial RoPE cos (channels 0-31 by x, 32-63 by y, 16 freqs theta 100 duplicated)",
            "sin": "[1, 1, 2520, 64] float16", "patch_mask": "[1, 2520] float16 (1 = patch, 0 = padding)"},
            "output": {"soft_tokens": "[1, 280, 512] float16; rows >= n_soft are padding"},
            "patch_size": 16, "pooling_kernel": 3, "cell_px": 48, "rope_theta": 100.0, "head_dim": 64,
            "resize": "aspect-preserving to multiples of 48 px within 2520 patches; antialiased bicubic on uint8 (torchvision-exact)",
            "image_prompt": None, "exif_orientation": "not applied (upstream parity); optional in the runners"},
        "image_sequence": "<bos> <|image> <|image|> x n_soft <image|> <eos>  (n_soft + 4 tokens, at most 284 -> bucket 288)",
        "audio": {**audio, "inputs": {
            "features": "[1, Tb, 128] float16 log-mel frames (10 ms hop, 128 HTK mel bands, ln(mel + 1e-3)), the first n_frames rows real, zero after",
            "frame_mask": "[1, Tb] float16 (1 for the first n_frames rows)",
            "half_mask": "[1, Tb/2] float16 (1 for the first ceil(n_frames/2) rows)",
            "token_mask": "[1, Tb/4] float16 (1 for the first ceil(ceil(n_frames/2)/2) rows = the soft-token count)"},
            "output": {"soft_tokens": "[1, Tb/4, 512] float16; rows >= n_tokens are padding"},
            "sample_rate": 16000, "max_seconds": 30.0, "max_samples": 480000, "frame_ms": 20, "hop_ms": 10, "fft": 512, "n_mels": 128,
            "frames_per_token": 4, "front_end": "python/eg2_audio_frontend.py, swift AudioFrontend (upstream Gemma4AudioFeatureExtractor defaults)",
            "resampler": "Kaiser-windowed sinc, pass band 0.913 / stop band 1.0 of the lower Nyquist, 125 dB (upstream: libsoxr HQ)",
            "audio_prompt": None, "long_audio": "the first 30 s (upstream); the runners also offer the mean of 30 s windows"},
        "audio_sequence": "<bos> <|audio> <|audio|> x n_tokens <audio|> <eos>  (n_tokens + 4 tokens, at most 754 -> bucket 1024)",
        "hidden_size": 512, "embedding_dim": 768, "matryoshka_dims": [768, 512, 256, 128],
        "bos_token_id": SPECIAL["bos"], "eos_token_id": SPECIAL["eos"], "pad_token_id": SPECIAL["pad"],
        "boi_token_id": SPECIAL["boi"], "image_token_id": SPECIAL["image"], "eoi_token_id": SPECIAL["eoi"],
        "video_token_id": SPECIAL["video"], "boa_token_id": SPECIAL["boa"], "audio_token_id": SPECIAL["audio"],
        "eoa_token_id": SPECIAL["eoa"], "prompts": PROMPTS,
        "tokenizer": "tokenizer/tokenizer.json (unchanged Gemma tokenizer; the runner adds <bos> and <eos>)",
    }
    (out / "config.json").write_text(json.dumps(cfg_json, indent=2) + "\n")
    import coreai_torch as ct
    from importlib.metadata import version as pkg_version

    tb = text["buckets"]
    attribution = {
        "name": "EmbeddingGemma-2-CoreAI",
        "version": "1.2",
        "base_model": args.hf_id,
        "base_model_author": "Gemma Team, Google DeepMind",
        "base_model_source": f"https://huggingface.co/{args.hf_id}",
        "base_model_license": "Apache-2.0",
        "base_model_terms": "Deployments must adhere to the Gemma Prohibited Use Policy: https://ai.google.dev/gemma/prohibited_use_policy",
        "base_model_file_sha256": {"model.safetensors": sha256_file(Path(args.model_dir) / "model.safetensors")},
        "port": {
            "copyright": "2026 Everpaper",
            "license": "Apache-2.0 (see LICENSE); attribution and modification notice in NOTICE",
            "date": args.date or time.strftime("%Y-%m-%d"),
            "format": "Apple Core AI .aimodel (macOS 27), fp16, static shapes",
            "exporter": {
                "coreai_torch": ct.__version__, "coreai_core": pkg_version("coreai-core"), "torch": torch.__version__.split("+")[0],
                "recipe": f"export/export.py --buckets {','.join(map(str, tb))} --qchunks {vision['qchunks']} "
                          f"--audio-buckets {','.join(map(str, audio['frame_buckets']))}",
                "global_attention": eg2_text.GLOBAL_ATTN_MODE, "gqa": eg2_text.GQA_MODE,
            },
        },
        "modifications": [
            "All three encoders of the multimodal checkpoint (text, vision, audio) are included, as separate assets.",
            "Architecture re-authored in plain PyTorch for Core AI export; weights are the original ones (bf16 rounded "
            "to fp16), with the scale folds and table cuts listed below.",
            "Audio: encoder (Gemma4AudioModel: subsampling convolutions, 12 conformer layers) and audio embedder re-authored "
            "in plain PyTorch; the chunked local attention is written in its banded form (each token attends to itself and "
            "the 11 previous tokens; the relative-position term is a per-layer constant), the q/k scales are folded into "
            "q_proj/k_proj and their clip bounds, the gradient-clipping clamps are dropped, the padded-key mask is an "
            "additive -10000, every RMSNorm uses the max-abs rescaling (fp16). Mathematically identical to the reference "
            "(fp32 soft-token cosine >= 0.99999998 on 213 clips).",
            f"Audio: fixed frame buckets {', '.join(map(str, audio['frame_buckets']))} (10 ms frames) as functions audio_<Tb>; "
            "inputs longer than 30 s are truncated by the host (the upstream default). The three masks are host inputs.",
            "Audio preprocessing (host): the upstream log-mel feature extractor re-implemented in numpy and Swift (bit-exact "
            "in numpy); resampling to 16 kHz with a Kaiser-windowed sinc instead of upstream's libsoxr (see README).",
            "Text: the sqrt(512) token-embedding scale and the 512^-0.5 per-layer-embedding scale are folded into the tables.",
            "Text: the four global (full-attention, single-kv-head) layers compute attention as a plain 4-D matmul with k/v "
            "broadcast over the heads; the sliding-window layers use a 5-D grouped matmul. Mathematically identical to "
            "the reference (fp32 max abs diff 3e-7).",
            "fp16 everywhere (Neural Engine): every RMSNorm and the final L2 normalization divide the row by its max-abs "
            "value before squaring (text activations reach ~2000, the vision residual ~1400); in the model code every "
            "division is expressed as rsqrt(x)*rsqrt(x) (the Neural Engine compiler rejects divide).",
            "Text: bidirectional attention masks (padding + sliding window 512) are built inside the graph from "
            "attention_mask; the additive mask value is -10000 (finite).",
            "Text: masked mean pooling, the 512->768 embedding projection and the L2 normalization run inside the graph; "
            "the output is the final sentence embedding.",
            "Text: an input mix x = embed(ids)*(1-mm_mask) + mm_embeds*mm_mask at the start of the graph, so that image "
            "soft tokens replace the <|image|> placeholder embeddings (soft tokens are not scaled); with zero soft tokens "
            "the graph is the v1.0 text encoder.",
            f"Text: fixed sequence buckets {', '.join(map(str, tb))} as functions mm_<S>; input longer than "
            f"{text['max_tokens']} tokens must be truncated or chunked by the caller (the original model accepts 8192).",
            "Vision: encoder (Gemma4VisionModel, 16 layers) and vision embedder, one fixed shape: the default budget of 280 "
            "soft tokens (2520 patches of 16x16). Patches are fed in cell-major order so that the 3x3 average pooling is a "
            "reshape + mean instead of a one-hot matmul (attention is permutation-equivariant; identical result).",
            "Vision: the patches are integer pixel values 0..255 (exact in fp16) instead of pixel/255: the graph computes "
            "P - 127.5 and the 2/255 scale is folded into the patch-embedding projection's weights, which equals upstream's "
            "projection of 2*(P/255) - 1 without rounding the inputs to fp16 (that rounding cost heavily upscaled tiny "
            "images up to 0.004 of cosine).",
            "Vision: the pooler's multiplication by sqrt(768) (fp32 upstream, overflows fp16) is dropped and the vision "
            "embedder's RMSNorm uses eps/768 instead: RMSNorm(x*s, eps) == RMSNorm(x, eps/s^2).",
            "Vision: the 2-D position-embedding tables are cut to the 840 rows reachable at this budget; the axial RoPE "
            "cos/sin tables and the key mask are inputs computed by the host; the additive key-mask value is -10000.",
            "Image preprocessing (host): the upstream aspect-preserving size rule and the torchvision-exact antialiased "
            "bicubic resize on uint8 re-implemented in numpy and Swift.",
        ],
    }
    if args.num_layers:
        attribution["port"]["exporter"]["num_layers"] = args.num_layers
    (out / "MODEL-ATTRIBUTION.json").write_text(json.dumps(attribution, indent=2) + "\n")
    log(f"bundle written: {out} ({dir_bytes(out) / 2**20:.1f} MiB)")
    log("DONE")


if __name__ == "__main__":
    main()