parakeet-rs 0.4.2

# Generate the burn backend's model code (src/burn_backend/generated/) from ONNX exports.
# Users never run this: the generated code is checked in. Run it again only when upgrading burn
# or adding a model.
#
# burn runs models written as Rust code, not .onnx files. For each ONNX graph of a model:
#   1. onnx2burn (burn-onnx's converter, same version as the `burn` dependency) writes the graph
#      as Rust model code, plus a .bpk file with the same weights in burn's own format.
#   2. The model code is trimmed to what parakeet-rs uses (no .bpk loading), and sizes that vary
#      between checkpoints of one architecture (e.g. the vocabulary) become parameters.
#   3. Every tensor in the .bpk is traced back to its ONNX source by content, giving a table that
#      parakeet-rs uses to load weights straight from the user's .onnx files, so users download
#      nothing new and no .bpk is ever shipped:
#        - model weights are looked up by a key that does not depend on the export: the
#          initializer's name, or for anonymous initializers ("onnx::MatMul_123") the nodes
#          that consume it;
#        - small graph constants (scalars) are embedded in the table.
#      Anything that cannot be traced is an error. The .bpk is then thrown away.
# No ONNX file is modified.
#
# Requirements:
#   cargo install burn-onnx --version 0.22.0 --locked     (provides onnx2burn)
#   pip install onnx numpy cbor2
#
# Usage:
#   python scripts/generate_burn.py <model> <model dir> [--onnx2burn path/to/onnx2burn]
# where <model> is one of the MODELS below, and <model dir> holds its reference export.
# Then check the result against ONNX Runtime with scripts/check_burn_parity.sh.

import argparse
import hashlib
import mmap
import os
import re
import struct
import subprocess
import tempfile

import cbor2
import numpy as np
import onnx
from onnx import numpy_helper

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "src", "burn_backend", "generated")
BURN_ONNX_VERSION = "0.22.0"
# Graph constants up to this many elements are embedded in the weight table.
MAX_EMBEDDED = 64


HEADER = "// Generated by scripts/generate_burn.py from {src} with burn-onnx {ver}. Do not edit.\n"


def digest(a):
    return hashlib.sha1(np.ascontiguousarray(a).tobytes()).hexdigest()


def rust_f32(x):
    """A Rust f32 literal for `x`, including infinities and NaN (used by attention masks)."""
    x = float(x)
    if np.isnan(x):
        return "f32::NAN"
    if np.isinf(x):
        return "f32::INFINITY" if x > 0 else "f32::NEG_INFINITY"
    return repr(x)


def weight_key(name, consumers):
    """Same rule as `weight_key` in src/burn_backend/onnx.rs."""
    if name.startswith("onnx::"):
        return "|".join(sorted(consumers.get(name, [])))
    return name


def read_bpk(path):
    # Memory-mapped: a .bpk holds every weight of the model (gigabytes for Cohere's encoder).
    with open(path, "rb") as f:
        raw = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    assert raw[:4] == b"NRUB", f"{path} is not a burnpack file"
    meta_len = struct.unpack("<I", raw[6:10])[0]
    entries = cbor2.loads(raw[10:10 + meta_len])["tensors"]
    return raw, meta_len, entries


def trace(onnx_path, bpk_path):
    """Map each .bpk tensor to (path, shape, source string for the Rust table)."""
    # Weights are read one at a time: a large model does not fit in memory twice.
    model = onnx.load(onnx_path, load_external_data=False)
    base_dir = os.path.dirname(onnx_path)
    protos = {t.name: t for t in model.graph.initializer}

    def init(name):
        return numpy_helper.to_array(protos[name], base_dir)
    graph_consts = {}
    for node in model.graph.node:
        if node.op_type == "Constant":
            for attr in node.attribute:
                if attr.name == "value":
                    graph_consts[node.output[0]] = numpy_helper.to_array(attr.t)
    consumers = {}
    for node in model.graph.node:
        for i, name in enumerate(node.input):
            if name in protos or name in graph_consts:
                consumers.setdefault(name, []).append(f"{node.name}#{i}")

    # content hash -> Rust source expression
    by_hash = {}
    for name in protos:
        by_hash.setdefault(digest(init(name)), f'Param("{weight_key(name, consumers)}", Copy)')
    for name, arr in graph_consts.items():
        if arr.size > MAX_EMBEDDED:
            # A model buffer exported as a constant (positional encoding tables). Read it from the
            # user's file; exports differ in the length of [1, N, D] tables, see CenterRows.
            transform = "CenterRows" if arr.ndim == 3 and arr.shape[0] == 1 else "Copy"
            by_hash.setdefault(digest(arr), f'Param("{weight_key(name, consumers)}", {transform})')
        elif arr.dtype == np.float32:
            values = ", ".join(rust_f32(x) for x in arr.ravel())
            by_hash.setdefault(digest(arr), f"ConstF32(&[{values}])")
        elif arr.dtype == np.int64:
            values = ", ".join(str(int(x)) for x in arr.ravel())
            by_hash.setdefault(digest(arr), f"ConstI64(&[{values}])")
    # LSTM: W/R are [1, 4H, I] (gates i, o, f, c); burn keeps one [I, H] matrix per gate. B is
    # [1, 8H] (W biases, then R biases); burn folds both into the input bias.
    bias_sums = []
    for node in model.graph.node:
        if node.op_type != "LSTM":
            continue
        h = next(a.i for a in node.attribute if a.name == "hidden_size")
        for name in node.input[1:3]:
            key = weight_key(name, consumers)
            for g in range(4):
                by_hash.setdefault(digest(init(name)[0][g * h:(g + 1) * h].T), f'Param("{key}", RowsT({g * h}, {(g + 1) * h}))')
        if len(node.input) > 3 and node.input[3]:
            name = node.input[3]
            b = init(name)[0]
            for g in range(4):
                expr = f'Param("{weight_key(name, consumers)}", BiasSum({g * h}, {(g + 1) * h}, {(g + 4) * h}, {(g + 5) * h}))'
                bias_sums.append((b[g * h:(g + 1) * h] + b[(g + 4) * h:(g + 5) * h], expr))

    raw, meta_len, entries = read_bpk(bpk_path)
    starts = sorted({((10 + meta_len + a - 1) // a) * a for a in (1, 8, 16, 32, 64, 256, 1024, 4096)})

    def payload(start, v):
        return raw[start + v["data_offsets"][0]:start + v["data_offsets"][1]]

    start = max(starts, key=lambda s: sum(hashlib.sha1(payload(s, v)).hexdigest() in by_hash for v in list(entries.values())[:64]))
    rows, missing = [], []
    for path, v in entries.items():
        data = payload(start, v)
        expr = by_hash.get(hashlib.sha1(data).hexdigest())
        if expr is None and v["dtype"] == "F32":
            arr = np.frombuffer(data, np.float32)
            if not arr.any():
                expr = "Zeros"
            else:
                expr = next((e for s, e in bias_sums if s.shape == arr.shape and np.allclose(s, arr, rtol=0, atol=1e-6)), None)
        if expr is None and np.prod(v["shape"], dtype=np.int64) <= MAX_EMBEDDED:
            # A small constant burn-onnx folded from graph constants (e.g. a Cast it evaluated):
            # not in the ONNX file, so embed burn's own value.
            kind, dtype = {"F32": ("ConstF32", np.float32), "I64": ("ConstI64", np.int64), "I32": ("ConstI32", np.int32)}.get(v["dtype"], (None, None))
            if kind:
                values = np.frombuffer(data, dtype)
                fmt = rust_f32 if dtype == np.float32 else (lambda x: str(int(x)))
                expr = f"{kind}(&[{', '.join(fmt(x) for x in values)}])"
        if expr is None:
            missing.append(path)
        else:
            rows.append((path, v["shape"], expr))
    if missing:
        raise SystemExit(f"{onnx_path}: {len(missing)} burn tensors have no ONNX source, e.g. {missing[:5]}")
    return rows


def trim_model_code(code):
    """Drop the .bpk loaders and the Default impl (it points at the .bpk path on the generating machine)."""
    code = re.sub(r"\nimpl Default for Model \{.*?\n\}\n", "\n", code, count=1, flags=re.S)
    code = re.sub(r"\nimpl Model \{\n    /// Load model weights from a burnpack file\..*?\n\}\n", "\n", code, count=1, flags=re.S)
    for line in ("use burn::tensor::Bytes;\n", "use burn_store::BurnpackStore;\n", "use burn_store::ModuleSnapshot;\n", "extern crate std;\n"):
        code = code.replace(line, "")
    assert "BurnpackStore" not in code and "from_file" not in code, "unexpected loader code left in generated model"
    return code


def tdt_decoder(code, rows):
    """The decoder's embedding and joint output sizes depend on the vocabulary: make them arguments."""
    shapes = {p: s for p, s, _ in rows}
    vocab, joint = shapes["constant1"][0], shapes["linear3.bias"][0]
    n = code.count(f"[{vocab}, 640]")
    assert n == 2, f"expected the [{vocab}, 640] embedding twice, found {n}"
    code = code.replace(f"[{vocab}, 640]", "[vocab_size, 640]")
    assert f"LinearConfig::new(640, {joint})" in code, "joint output layer not found"
    code = code.replace(f"LinearConfig::new(640, {joint})", "LinearConfig::new(640, joint_size)")
    code = code.replace("pub fn new(device: &Device) -> Self {", "pub fn new(device: &Device, vocab_size: usize, joint_size: usize) -> Self {", 1)
    assert "vocab_size: usize" in code
    # vocabulary-sized tensors take their shape from the user's file
    rows = [(p, None if p in ("constant1", "linear3.weight", "linear3.bias") else s, e) for p, s, e in rows]
    return code, rows


def static_shapes(model):
    """Fully static tensor shapes from ONNX shape inference, by name.

    Weights are replaced by typed graph inputs first: inference serializes the model, and a
    multi-GB model does not fit in one protobuf.
    """
    light = onnx.ModelProto()
    light.CopyFrom(model)
    graph = light.graph
    graph.input.extend(
        onnx.helper.make_tensor_value_info(t.name, t.data_type, list(t.dims)) for t in graph.initializer
    )
    del graph.initializer[:]
    inferred = onnx.shape_inference.infer_shapes(light)
    shapes = {}
    for v in list(inferred.graph.value_info) + list(inferred.graph.input) + list(inferred.graph.output):
        dims = v.type.tensor_type.shape.dim
        if dims and all(d.HasField("dim_value") for d in dims):
            shapes[v.name] = [d.dim_value for d in dims]
    return shapes


# Inputs that burn-onnx needs on the host: (op, input index).
SHAPE_INPUTS = {("Pad", 1), ("Expand", 1)}


def fold_constant_shapes(model):
    """Pad amounts and Expand sizes computed from constants only -> one constant.

    PyTorch exports such values as small ConstantOfShape/Concat/Slice/Transpose/Reshape chains.
    burn-onnx 0.22 evaluates a chain on the device and reads the result back to the host on every
    call; each read waits for the GPU queue (and breaks graph capture). The chain depends on
    constants only, so it is evaluated here once with ONNX's reference evaluator.
    """
    from onnx.reference import ReferenceEvaluator

    graph = model.graph
    produced = {o: n for n in graph.node for o in n.output}
    inputs = {i.name for i in graph.input}
    shapes = static_shapes(model)
    small_inits = {t.name: t for t in graph.initializer if np.prod(t.dims, dtype=np.int64) <= MAX_EMBEDDED}

    def constant_chain(name, chain):
        """Nodes computing `name` if it depends on constants only, else None."""
        if name in small_inits:
            chain[name] = small_inits[name]
            return chain
        if name in inputs:
            return None
        node = produced.get(name)
        if node is None:
            return None
        if node.op_type == "Constant":
            chain[node.output[0]] = node
            return chain
        if node.op_type == "Shape":
            # Shape of a tensor whose shape is static: a constant, even if its values are not.
            dims = shapes.get(node.input[0])
            if dims is None:
                return None
            attrs = {a.name: a.i for a in node.attribute}
            value = np.array(dims[attrs.get("start", 0):attrs.get("end", len(dims))], np.int64)
            chain[node.output[0]] = onnx.helper.make_node("Constant", [], [node.output[0]], value=numpy_helper.from_array(value))
            return chain
        for i in node.input:
            if i and constant_chain(i, chain) is None:
                return None
        chain[node.output[0]] = node
        return chain

    folded = 0
    new_constants = []
    for node in graph.node:
        index = next((i for op, i in SHAPE_INPUTS if op == node.op_type), None)
        if index is None or len(node.input) <= index or not node.input[index]:
            continue
        source = produced.get(node.input[index])
        if source is None or source.op_type == "Constant":
            # Already a constant (an initializer or a Constant node).
            continue
        chain = constant_chain(node.input[index], {})
        if chain is None:
            continue
        nodes = [n for n in chain.values() if isinstance(n, onnx.NodeProto)]
        inits = [t for t in chain.values() if isinstance(t, onnx.TensorProto)]
        # Folded shapes are new Constant nodes outside the graph: they go first, then the graph's
        # own nodes of the chain in graph order.
        in_graph = [n for n in graph.node if any(n is m for m in nodes)]
        order = [n for n in nodes if not any(n is m for m in in_graph)] + in_graph
        target = node.input[index]
        sub = onnx.helper.make_model(
            onnx.helper.make_graph(order, "shape", [], [onnx.helper.make_tensor_value_info(target, onnx.TensorProto.INT64, None)], inits),
            opset_imports=model.opset_import,
        )
        (value,) = ReferenceEvaluator(sub).run(None, {})
        name = f"{target}/folded"
        new_constants.append(onnx.helper.make_node("Constant", [], [name], value=numpy_helper.from_array(np.asarray(value, np.int64), name)))
        node.input[index] = name
        folded += 1
    for c in reversed(new_constants):
        graph.node.insert(0, c)
    print(f"  folded {folded} constant pad amounts / expand sizes")
    return model


def expand_shape_to_abs(model):
    """Where(s == -1, 1, s) feeding Expand -> |s|.

    PyTorch exports x.expand(..., -1, ...) by replacing the -1 entries of the size vector with 1.
    burn-onnx 0.22 runs that Where on the device and reads the shape back on every call. A size
    vector holds -1 or non-negative sizes only, so |s| is the same vector, and burn-onnx keeps
    it on the host. Matched strictly: Where(Equal(s, Mul(ones, -1)), ones, s) with ones a
    ConstantOfShape of value 1.
    """
    graph = model.graph
    produced = {o: n for n in graph.node for o in n.output}

    def const_value(name):
        n = produced.get(name)
        if n is None:
            return None
        if n.op_type == "Constant":
            return numpy_helper.to_array(n.attribute[0].t)
        if n.op_type == "ConstantOfShape":
            return numpy_helper.to_array(n.attribute[0].t) if n.attribute else np.array([0.0])
        return None

    rewritten = 0
    for node in graph.node:
        if node.op_type != "Expand":
            continue
        where = produced.get(node.input[1])
        if where is None or where.op_type != "Where":
            continue
        cond, ones, s = where.input
        equal, ones_node = produced.get(cond), produced.get(ones)
        if equal is None or equal.op_type != "Equal" or ones_node is None or ones_node.op_type != "ConstantOfShape":
            continue
        if list(const_value(ones).ravel()) != [1] or equal.input[0] != s:
            continue
        mul = produced.get(equal.input[1])
        if mul is None or mul.op_type != "Mul" or mul.input[0] != ones:
            continue
        if list(np.asarray(const_value(mul.input[1])).ravel()) != [-1]:
            continue
        # |s| as Max(s, -s): burn-onnx 0.22 keeps Mul and Max of shape values on the host but
        # generates uncompilable code for Abs of one.
        out = f"{node.input[1]}/abs"
        minus_one = f"{out}/minus_one"
        at = list(graph.node).index(node)
        graph.node.insert(at, onnx.helper.make_node("Max", [s, f"{out}/neg"], [out], name=f"{where.name}/abs"))
        graph.node.insert(at, onnx.helper.make_node("Mul", [s, minus_one], [f"{out}/neg"], name=f"{where.name}/neg"))
        graph.node.insert(at, onnx.helper.make_node("Constant", [], [minus_one], value=numpy_helper.from_array(np.array(-1, np.int64))))
        node.input[1] = out
        rewritten += 1
    print(f"  rewrote {rewritten} expand size vectors as |s|")
    return model


def default_patches(model):
    return expand_shape_to_abs(fold_constant_shapes(model))


def shape_scatter_to_concat(model):
    """ScatterElements that overwrite one entry of a shape vector -> the equivalent Concat.

    PyTorch exports index_put/expand bookkeeping as ScatterElements(data=Shape(x), indices=[k],
    updates=u). burn-onnx 0.22 cannot type a scatter into a shape value; Concat(data[:k], u,
    data[k+1:]) computes the same vector and is supported. Only the code-generation copy of the
    graph is rewritten; no weights are involved.
    """
    graph = model.graph
    produced = {o: n for n in graph.node for o in n.output}
    const = {}
    for n in graph.node:
        if n.op_type == "Constant":
            const[n.output[0]] = numpy_helper.to_array(n.attribute[0].t)
    nodes, rewritten = [], 0
    for n in graph.node:
        data = produced.get(n.input[0]) if n.op_type == "ScatterElements" else None
        idx = const.get(n.input[1]) if data is not None else None
        if data is None or data.op_type != "Shape" or idx is None or idx.size != 1:
            nodes.append(n)
            continue
        k, out = int(idx.ravel()[0]), n.output[0]
        assert k >= 0, f"{n.name}: negative scatter index"
        parts = []
        if k > 0:
            nodes.append(onnx.helper.make_node("Slice", [n.input[0], f"{out}/k0", f"{out}/k"], [f"{out}/head"], name=f"{n.name}/head"))
            parts.append(f"{out}/head")
        parts.append(n.input[2])
        nodes.append(onnx.helper.make_node("Slice", [n.input[0], f"{out}/k1", f"{out}/end"], [f"{out}/tail"], name=f"{n.name}/tail"))
        parts.append(f"{out}/tail")
        for name, value in ((f"{out}/k0", 0), (f"{out}/k", k), (f"{out}/k1", k + 1), (f"{out}/end", 2**62)):
            nodes.insert(0, onnx.helper.make_node("Constant", [], [name], value=numpy_helper.from_array(np.array([value], np.int64))))
        nodes.append(onnx.helper.make_node("Concat", parts, [out], name=n.name, axis=0))
        rewritten += 1
    del graph.node[:]
    graph.node.extend(nodes)
    print(f"  rewrote {rewritten} shape scatters as Concat")
    return model


def declare_opset_19(model):
    """Declare opset 19 on an opset-17 graph, unchanged otherwise.

    burn-onnx 0.22 only reads AveragePool's ceil_mode from opset 19, although ONNX has had it
    since opset 10. Between 17 and 19 the ops used here only gained types and optional
    attributes, except Split, which then needs explicit sizes; anything else is refused.
    """
    from onnx import defs
    additive = {"AveragePool", "Cast", "Constant", "Equal", "Reshape", "ScatterElements", "Shape", "Split"}
    (opset,) = [o for o in model.opset_import if o.domain in ("", "ai.onnx")]
    assert opset.version == 17, f"expected opset 17, got {opset.version}"
    for op in {n.op_type for n in model.graph.node}:
        changed = defs.get_schema(op, 17).since_version != defs.get_schema(op, 19).since_version
        assert not changed or op in additive, f"{op} changes meaning between opset 17 and 19"
    assert all(len(n.input) > 1 and n.input[1] for n in model.graph.node if n.op_type == "Split"), "Split without sizes"
    opset.version = 19
    return model


def sortformer_patches(model):
    return declare_opset_19(shape_scatter_to_concat(default_patches(model)))


def inline_if(model, branch):
    """Replace every top-level If with the nodes of one branch ("then_branch" or "else_branch").

    Cohere's merged decoder switches on "is this the first step" (empty encoder cache). parakeet-rs
    runs the two cases as two graphs, so each is a plain graph without If, and the first one
    needs no empty past tensors (burn cannot handle zero-length tensors).
    """
    graph = model.graph
    nodes = []
    for node in graph.node:
        if node.op_type != "If":
            nodes.append(node)
            continue
        (body,) = [a.g for a in node.attribute if a.name == branch]
        graph.initializer.extend(body.initializer)
        nodes.extend(body.node)
        for inner, outer in zip(body.output, node.output):
            if inner.name != outer:
                nodes.append(onnx.helper.make_node("Identity", [inner.name], [outer], name=f"{node.name}/{outer}"))
    del graph.node[:]
    graph.node.extend(nodes)
    return model


def prune(model, drop_outputs=()):
    """Drop the given graph outputs, then every node, input and initializer they no longer need."""
    graph = model.graph
    keep = [o for o in graph.output if o.name not in drop_outputs]
    del graph.output[:]
    graph.output.extend(keep)
    needed = {o.name for o in graph.output}
    nodes = []
    for node in reversed(graph.node):
        if any(o in needed for o in node.output):
            nodes.append(node)
            needed.update(i for i in node.input if i)
    del graph.node[:]
    graph.node.extend(reversed(nodes))
    for field in (graph.input, graph.initializer):
        kept = [x for x in field if x.name in needed]
        del field[:]
        field.extend(kept)
    return model


def ms_attention_to_onnx(model, causal_without_past=False):
    """com.microsoft MultiHeadAttention / GroupQueryAttention -> the standard Attention (opset 23).

    burn-onnx has no com.microsoft ops. In the Cohere export they are plain attention:
    - MultiHeadAttention: Q, K, V and an optional additive attention_bias, nothing else. Q, K, V
      are [B, S, heads * dim], or K and V already [B, heads, L, dim] (decoder cross-attention).
    - GroupQueryAttention: as many KV heads as query heads, no rotary, no softcap, no local
      window, so it is self-attention over past + new keys, with an additive bias and the causal
      mask. The causal mask only matters when several new tokens attend to each other: true for
      the first step (no past), where it is the plain square mask (is_causal); a later step must
      feed one token (checked in Rust), which may attend to everything.
    """
    graph = model.graph
    produced = {o: n for n in graph.node for o in n.output}
    inputs = {i.name: i for i in graph.input}

    def rank(name):
        if name in inputs:
            return len(inputs[name].type.tensor_type.shape.dim)
        node = produced[name]
        if node.op_type == "Identity":
            return rank(node.input[0])
        if node.op_type == "Transpose":
            return len(next(a.ints for a in node.attribute if a.name == "perm"))
        return 3

    out = []

    def add(op, ins, outs, **attrs):
        out.append(onnx.helper.make_node(op, ins, outs, name=outs[0], **attrs))
        return outs[0]

    def const(name, values):
        graph.initializer.append(numpy_helper.from_array(np.array(values, np.int64), name))
        return name

    def to_4d(x, heads, prefix):
        r = add("Reshape", [x, const(f"{prefix}/split_shape", [0, 0, heads, -1])], [f"{prefix}/split"])
        return add("Transpose", [r], [f"{prefix}/bnsh"], perm=[0, 2, 1, 3])

    def to_3d(x, y, prefix):
        t = add("Transpose", [x], [f"{prefix}/bsnh"], perm=[0, 2, 1, 3])
        add("Reshape", [t, const(f"{prefix}/merge_shape", [0, 0, -1])], [y])

    for node in graph.node:
        if node.domain != "com.microsoft":
            out.append(node)
            continue
        attrs = {a.name: onnx.helper.get_attribute_value(a) for a in node.attribute}
        heads, scale = attrs["num_heads"], attrs["scale"]
        ins = list(node.input) + [""] * 11
        p = node.name
        if node.op_type == "MultiHeadAttention":
            q, k, v, bias = ins[0], ins[1], ins[2], ins[5]
            assert not any(ins[3:5]) and not any(ins[6:8]), f"{p}: unsupported MultiHeadAttention inputs"
            if rank(k) == 4:
                y4 = add("Attention", [to_4d(q, heads, f"{p}/q"), k, v, bias], [f"{p}/y4"], scale=scale)
                to_3d(y4, node.output[0], p)
            else:
                add("Attention", [q, k, v, bias], [node.output[0]], scale=scale,
                    q_num_heads=heads, kv_num_heads=heads)
        elif node.op_type == "GroupQueryAttention":
            assert attrs["kv_num_heads"] == heads and attrs["do_rotary"] == 0, f"{p}: not plain attention"
            assert attrs.get("softcap", 0.0) == 0.0 and attrs.get("local_window_size", -1) == -1, f"{p}: windowed or softcapped"
            assert not ins[7] and not ins[8], f"{p}: rotary inputs"
            q, k, v, past_k, past_v, bias = ins[0], ins[1], ins[2], ins[3], ins[4], ins[10]
            y, present_k, present_v = node.output
            if causal_without_past:
                k4, v4 = to_4d(k, heads, f"{p}/k"), to_4d(v, heads, f"{p}/v")
                y4 = add("Attention", [to_4d(q, heads, f"{p}/q"), k4, v4, bias], [f"{p}/y4"], scale=scale, is_causal=1)
                to_3d(y4, y, p)
                add("Identity", [k4], [present_k])
                add("Identity", [v4], [present_v])
            else:
                add("Attention", [q, k, v, bias, past_k, past_v], [y, present_k, present_v], scale=scale,
                    q_num_heads=heads, kv_num_heads=heads)
        else:
            raise SystemExit(f"{p}: unsupported com.microsoft op {node.op_type}")
    del graph.node[:]
    graph.node.extend(out)
    opsets = [o for o in model.opset_import if o.domain in ("", "ai.onnx")]
    del model.opset_import[:]
    model.opset_import.extend(opsets)
    return declare_opset_23(model)


def declare_opset_23(model):
    """Declare opset 23 (for Attention) on an opset-21 graph. Between 21 and 23 the other ops
    used here only gained types; anything else is refused."""
    from onnx import defs
    (opset,) = [o for o in model.opset_import if o.domain in ("", "ai.onnx")]
    assert opset.version == 21, f"expected opset 21, got {opset.version}"
    for op in {n.op_type for n in model.graph.node} - {"Attention"}:
        a, b = defs.get_schema(op, 21), defs.get_schema(op, 23)
        same = sorted(a.attributes) == sorted(b.attributes) and len(a.inputs) == len(b.inputs)
        assert same, f"{op} changes between opset 21 and 23"
    opset.version = 23
    return model


def unshare_small_constants(model):
    """Give every use of a small initializer its own copy with a plain name; values unchanged.

    burn-onnx 0.22 turns ONNX names into Rust identifiers, and Cohere's constant names collide
    there ("/model/constants/INT64/[-1]" and ".../[1]"), so one value silently replaces the
    other (a Range step or a reshape size -1 became 1).
    """
    graph = model.graph
    small = {t.name: t for t in graph.initializer if np.prod(t.dims, dtype=np.int64) <= MAX_EMBEDDED}
    copies = []
    for node in graph.node:
        for k, name in enumerate(node.input):
            if name not in small:
                continue
            copy = onnx.TensorProto()
            copy.CopyFrom(small[name])
            copy.name = f"small_const_{len(copies)}"
            copies.append(copy)
            node.input[k] = copy.name
    kept = [t for t in graph.initializer if t.name not in small]
    del graph.initializer[:]
    graph.initializer.extend(kept + copies)
    return model


def cohere_encoder_patches(model):
    return unshare_small_constants(prune(default_patches(ms_attention_to_onnx(model))))


def cohere_first_step_patches(model):
    """The decoder's first step: the prompt, no past, computing the cross-attention cache."""
    model = inline_if(model, "then_branch")
    model = ms_attention_to_onnx(model, causal_without_past=True)
    return unshare_small_constants(default_patches(prune(model)))


def cohere_next_step_patches(model):
    """The decoder's later steps: one token, reusing the cross-attention cache from the first step."""
    model = inline_if(model, "else_branch")
    model = ms_attention_to_onnx(model)
    encoder_cache = [o.name for o in model.graph.output if o.name.endswith((".encoder.key", ".encoder.value"))]
    return unshare_small_constants(default_patches(prune(model, drop_outputs=encoder_cache)))


def graph(file, name, post=None, patch=None, simplify=True):
    """One ONNX graph of a model: its file in the model dir and the generated module name.

    post:     edits the generated code and weight table (e.g. vocabulary-sized parameters)
    patch:    rewrites the reference graph into an equivalent one burn-onnx can convert
    simplify: burn-onnx's graph simplification. Its 0.22 simplifier drops the reshapes around
              the relative-position shift (pad -> reshape -> slice -> reshape) in the
              Transformers-style export, giving wrong shapes; such graphs turn it off.
    """
    return dict(file=file, name=name, post=post, patch=patch or default_patches, simplify=simplify)


# model -> (reference export, its graphs)
MODELS = {
    "tdt": (
        "istupakov/parakeet-tdt-0.6b-v3-onnx",
        [graph("encoder-model.onnx", "tdt_encoder"), graph("decoder_joint-model.onnx", "tdt_decoder_joint", post=tdt_decoder)],
    ),
    "ctc": (
        "onnx-community/parakeet-ctc-0.6b-ONNX (onnx/)",
        [graph("model.onnx", "ctc", simplify=False)],
    ),
    "nemotron": (
        "altunenes/parakeet-rs nemotron-speech-streaming-en-0.6b",
        [graph("encoder.onnx", "nemotron_encoder"), graph("decoder_joint.onnx", "nemotron_decoder_joint")],
    ),
    "nemotron_multi": (
        "altunenes/parakeet-rs nemotron-3.5-asr-streaming-0.6b-onnx",
        [graph("encoder.onnx", "nemotron_multi_encoder"), graph("decoder_joint.onnx", "nemotron_multi_decoder_joint")],
    ),
    "eou": (
        "altunenes/parakeet-rs realtime_eou_120m-v1-onnx",
        [graph("encoder.onnx", "eou_encoder"), graph("decoder_joint.onnx", "eou_decoder_joint")],
    ),
    "unified": (
        "bobNight/parakeet-unified-en-0.6b-onnx (fp32)",
        [graph("encoder.onnx", "unified_encoder"), graph("decoder_joint.onnx", "unified_decoder_joint")],
    ),
    "multitalker": (
        "multitalker-parakeet-streaming-0.6b-v1 fp32 export (scripts/export_multitalker.py)",
        [graph("encoder.onnx", "multitalker_encoder"), graph("decoder_joint.onnx", "multitalker_decoder_joint")],
    ),
    "sortformer": (
        "altunenes/parakeet-rs nemotron-3-diarization (scripts/export_diar_sortformer.py)",
        [graph("nemotron3_diar_v3.onnx", "sortformer", patch=sortformer_patches, simplify=False)],
    ),
    "cohere": (
        "onnx-community/cohere-transcribe-03-2026-ONNX (onnx/, fp32)",
        [
            graph("encoder_model.onnx", "cohere_encoder", patch=cohere_encoder_patches),
            graph("decoder_model_merged.onnx", "cohere_decoder_first", patch=cohere_first_step_patches),
            graph("decoder_model_merged.onnx", "cohere_decoder_next", patch=cohere_next_step_patches),
        ],
    ),
}


def generate(onnx_path, out_name, post, patch, simplify, tmp, onnx2burn):
    name = os.path.splitext(os.path.basename(onnx_path))[0]
    source = onnx_path
    if patch:
        # Rewrite a copy of the graph for code generation; weights are still traced to the
        # original file. The copy sits next to the original (removed below), so it reads the same
        # external data files: no weights are loaded or copied, and ONNX tools refuse data files
        # outside the model's folder.
        source = os.path.join(os.path.dirname(onnx_path), f"{name}.burn-patched.onnx")
        onnx.save(patch(onnx.load(onnx_path, load_external_data=False)), source)
    flags = ["--no-development"] + ([] if simplify else ["--no-simplify"])
    try:
        run = subprocess.run([onnx2burn, *flags, source, os.path.join(tmp, name)],
                             stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, env={**os.environ, "RUST_LOG": "error"})
    finally:
        if source != onnx_path:
            os.remove(source)
    if run.returncode != 0:
        reason = [line for line in run.stderr.splitlines() if "Failed" in line or "panicked" in line]
        raise SystemExit(f"onnx2burn failed on {onnx_path}:\n" + "\n".join(reason[-2:] or run.stderr.splitlines()[-5:]))
    generated = os.path.join(tmp, name, os.path.splitext(os.path.basename(source))[0])
    code = open(f"{generated}.rs").read()
    # Drop onnx2burn's own header: it holds a local path.
    code = trim_model_code(code.split("\n", 1)[1])
    rows = trace(onnx_path, f"{generated}.bpk")
    if post:
        code, rows = post(code, rows)

    header = HEADER.format(src=os.path.basename(onnx_path), ver=BURN_ONNX_VERSION)
    with open(os.path.join(OUT, f"{out_name}.rs"), "w") as f:
        f.write(header + code)
    with open(os.path.join(OUT, f"{out_name}_weights.rs"), "w") as f:
        f.write(header)
        f.write("// Where each burn parameter's data comes from; see src/burn_backend/onnx.rs.\n")
        f.write("use crate::burn_backend::onnx::{Source::*, Transform::*, Weight};\n\n")
        f.write("pub(crate) const WEIGHTS: &[Weight] = &[\n")
        for path, shape, expr in rows:
            shape_expr = "None" if shape is None else f"Some(&{list(shape)})"
            f.write(f'    Weight {{ path: "{path}", shape: {shape_expr}, source: {expr} }},\n')
        f.write("];\n")
    kinds = {}
    for _, _, e in rows:
        k = re.match(r'(Param\("[^"]*", )?(\w+)', e).group(2)
        kinds[k] = kinds.get(k, 0) + 1
    print(f"{out_name}: {len(rows)} tensors {kinds}")


def main():
    parser = argparse.ArgumentParser(description="Generate burn model code for parakeet-rs.")
    parser.add_argument("model", choices=sorted(MODELS))
    parser.add_argument("model_dir", help="Folder with the model's reference ONNX export")
    parser.add_argument("--onnx2burn", default="onnx2burn", help=f"onnx2burn binary (burn-onnx {BURN_ONNX_VERSION})")
    args = parser.parse_args()

    version = subprocess.run([args.onnx2burn, "--version"], capture_output=True, text=True).stdout.strip()
    if BURN_ONNX_VERSION not in version:
        raise SystemExit(f"need onnx2burn {BURN_ONNX_VERSION}, found {version!r}")

    reference, graphs = MODELS[args.model]
    print(f"{args.model}: reference export {reference}")
    os.makedirs(OUT, exist_ok=True)
    with tempfile.TemporaryDirectory() as tmp:
        for g in graphs:
            generate(os.path.join(args.model_dir, g["file"]), g["name"], g["post"], g["patch"], g["simplify"], tmp, args.onnx2burn)


if __name__ == "__main__":
    main()