code/tt_frnet/reference/weights.py

# SPDX-License-Identifier: Apache-2.0
"""The deployed ``frnet_<variant>.onnx`` of AutowareFoundation/lidar_frnet@v2.0 read as DATA, the way the TT port
reads it: the vendored ``ttaw.weights.OnnxWeights`` (initializers addressed by the node that consumes them; nothing
is executed). 120 of the 186 initializer *names* differ between the OT128 and QT128 exports (``onnx::Conv_1487`` vs
``onnx::Conv_1485``, S:frnet:98-101), node names do not, so every weight is found from its node.

Where BatchNorm is folded (float64, one final rounding to float32; ``ttaw.weights.fold_bn_linear``):

- every point layer ``MatMul (no bias) -> BatchNormalization(eps 1e-3) -> Relu`` (``ffe_layers.0-2``,
  ``point_stem``, ``point_fusion_layers.0-3``, ``point_fuse_layer1-2``, ``mlps.0-3``) becomes one
  ``y = relu(x @ W' + b')`` -- exactly the port's ``ttnn.linear(..., bias, activation="relu")``;
- ``pre_norm`` (BatchNorm1d over the 8 decorated features) is kept as ``(scale, shift)`` float32 (``x * scale +
  shift``): the port applies it on the host with the feature decoration (PLAN.md 2.8), so the first device layer is
  a plain linear;
- the 53 convs already carry the exporter's BN fold (Conv with bias, S:frnet:245); the three Gemms
  (``ffe_layers.3``, ``compression_layers.0``, ``conv_seg``) have their own bias.

The structure is recovered from the graph, never re-typed: stage count, BasicBlocks per stage, strides, the
frustum grid of every frustum2pixel (the ``ConstantOfShape`` shape of its ``ScatterND``), every activation (the
conv's consumer), and every wiring assumption of the reference model (concat orders, skip connections, the
decoder's resize inputs, the head's residual skips, the Softmax axis) is asserted on the graph
(:func:`load_weights` raises ``ValueError`` when one does not hold). Each layer records the ONNX tensor that holds
its output, so ``reference/ort.py`` exposes the same taps from ONNX Runtime.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np

from ..ttaw.weights import BatchNorm, OnnxWeights, file_sha256, fold_bn_linear
from .config import BLOCKS, NUM_CLASSES, ONNX_SHA256, STRIDES, variant_files

__all__ = ["Linear", "Conv", "Stage", "FRNetWeights", "load_weights", "check_variant_grid", "MODULE_TAPS",
           "OPT_IN_TAPS"]

VE = "/model/voxel_encoder"
BB = "/model/backbone"


@dataclass(frozen=True)
class Linear:
    """``y = act(x @ weight + bias)``; ``weight`` (in, out) float32 (BN folded where the graph has one)."""

    name: str
    weight: np.ndarray
    bias: np.ndarray
    relu: bool
    onnx_node: str
    onnx_output: str                 # ONNX tensor equal to this layer's output (after the ReLU when there is one)
    bn_node: Optional[str] = None

    @property
    def in_features(self) -> int:
        return int(self.weight.shape[0])

    @property
    def out_features(self) -> int:
        return int(self.weight.shape[1])


@dataclass(frozen=True)
class Conv:
    """``y = act(conv2d(x, weight, bias, stride, padding))`` (NCHW, BN already folded by the exporter)."""

    name: str
    weight: np.ndarray               # (Cout, Cin, k, k) float32
    bias: np.ndarray                 # (Cout,) float32
    stride: int
    padding: int
    activation: Optional[str]        # "hardswish" | "sigmoid" | None (a residual add follows)
    onnx_node: str
    onnx_output: str                 # ONNX tensor equal to the output (after the activation when there is one)

    @property
    def kernel(self) -> int:
        return int(self.weight.shape[-1])

    @property
    def in_channels(self) -> int:
        return int(self.weight.shape[1])

    @property
    def out_channels(self) -> int:
        return int(self.weight.shape[0])


@dataclass(frozen=True)
class Stage:
    """One backbone stage ``layer<index>`` (1-based) plus its point / pixel fusion (``*_layers.<index - 1>``)."""

    index: int
    stride: int                      # cumulative stride of the stage output (1, 2, 4, 8)
    blocks: Tuple[Tuple[Conv, Conv, Optional[Conv]], ...]   # (conv1 + HardSwish, conv2, downsample) per BasicBlock
    block_outputs: Tuple[str, ...]   # ONNX tensor of each block output (after add + HardSwish)
    point_fusion: Linear             # [pixel2point(x_s), p] 256 -> 128 + BN + ReLU
    pixel_fusion: Conv               # [frustum2pixel(segmax(p)), x_s] 256 -> 128 + HardSwish
    attention: Tuple[Conv, Conv]     # 128 -> 128 + HardSwish, 128 -> 128 + Sigmoid
    segmax_output: str               # ONNX tensor: per-voxel max of point_fusion's output (M, 128)
    output: str                      # ONNX tensor: x = pixel_fusion * attention + x_s


@dataclass(frozen=True)
class FRNetWeights:
    variant: str
    frustum_hw: Tuple[int, int]
    strides: Tuple[int, ...]
    pre_norm: BatchNorm
    pre_norm_scale: np.ndarray       # (8,) float32: pre_norm(x) == x * scale + shift
    pre_norm_shift: np.ndarray
    ffe: Tuple[Linear, ...]          # 8 -> 64 -> 128 -> 256 (ReLU, BN folded), 256 -> 256 (Gemm, no activation)
    compression: Linear              # 256 -> 16 + ReLU (on the per-voxel max)
    stem: Tuple[Conv, ...]           # 16 -> 64 -> 128 -> 128, each + HardSwish
    point_stem: Linear               # [pixel2point(stem), pf3] 384 -> 128 + BN + ReLU
    point_stem_segmax: str           # ONNX tensor of the per-voxel max of point_stem (M, 128)
    fusion_stem: Conv                # [frustum2pixel(segmax(p0)), stem] 256 -> 128 + HardSwish
    stages: Tuple[Stage, ...]
    fuse_layer1: Conv                # concat of the 5 up-sampled maps 640 -> 256 + HardSwish
    fuse_layer2: Conv                # 256 -> 128 + HardSwish
    point_fuse: Tuple[Linear, Linear]   # 640 -> 256, 256 -> 128 (+ BN + ReLU)
    head_mlps: Tuple[Linear, ...]    # 128 -> 128, 128 -> 256, 256 -> 128, 128 -> 64 (+ BN + ReLU), then + skip
    conv_seg: Linear                 # 64 -> 27
    onnx_taps: Dict[str, str]        # tap name -> ONNX tensor (module and opt-in taps)
    onnx_path: str = ""
    onnx_sha256: str = ""
    raw_parameters: int = 0          # learned values in the ONNX initializers
    folded_parameters: int = field(default=0)   # values after the BN folds (what the device holds)

    @property
    def convs(self) -> List[Conv]:
        out = list(self.stem) + [self.fusion_stem]
        for st in self.stages:
            for c1, c2, ds in st.blocks:
                out += [c1, c2] + ([ds] if ds is not None else [])
            out += [st.pixel_fusion, *st.attention]
        return out + [self.fuse_layer1, self.fuse_layer2]

    @property
    def linears(self) -> List[Linear]:
        return [*self.ffe, self.compression, self.point_stem, *(s.point_fusion for s in self.stages),
                *self.point_fuse, *self.head_mlps, self.conv_seg]

    @property
    def num_classes(self) -> int:
        return self.conv_seg.out_features


# Module taps of the reference (tests/test_reference_cpu.py compares every one with ORT; the TT port's per-module
# gates use the same names). 1-based stage numbers: stage l = ONNX layer<l> + *_layers.<l - 1>.
MODULE_TAPS: Tuple[str, ...] = (
    "encoder.input", "encoder.pf0", "encoder.pf1", "encoder.pf2", "encoder.pf3", "encoder.voxel_max",
    "encoder.voxel_feats", "backbone.frustum", "backbone.stem", "backbone.point_stem", "backbone.fusion_stem",
    *(f"backbone.{kind}{i}" for i in range(1, 5)
      for kind in ("layer", "point_fusion", "pixel_fusion", "attention", "stage")),
    "decoder.fuse1", "decoder.fuse2", "decoder.point_fuse1", "decoder.point_fuse2",
    *(f"head.mlp{i}" for i in range(4)), "head.logits", "pred_probs",
)
# Opt-in taps (TapRegistry include patterns): every BasicBlock output, the per-voxel max of every point branch
# (the K1 / segmented-max sites), the five up-sampled decoder inputs.
OPT_IN_TAPS: Tuple[str, ...] = (
    "backbone.layer*.block*", "backbone.segmax*", "decoder.up*",
)


def _fail(msg: str) -> None:
    raise ValueError(f"frnet ONNX structure: {msg}")


def _single(w: OnnxWeights, tensor: str, op_type: str) -> Any:
    try:
        return w.consumer_of(tensor, op_type)
    except KeyError as e:
        _fail(str(e))


def _point_layer(w: OnnxWeights, matmul: str, name: str) -> Linear:
    mm = w.node(matmul)
    if mm.op_type != "MatMul" or not w.has(mm.inputs[1]) or w.has(mm.inputs[0]):
        _fail(f"{matmul} is not x @ W with a constant W")
    bn = _single(w, mm.outputs[0], "BatchNormalization")
    relu = _single(w, bn.outputs[0], "Relu")
    wf, bf = fold_bn_linear(w.matmul_weight(matmul), None, w.batchnorm(bn.name), layout="in_out")
    return Linear(name, wf, bf, True, mm.name, relu.outputs[0], bn.name)


def _gemm(w: OnnxWeights, node: str, name: str) -> Linear:
    g = w.gemm(node)
    if g.trans_a or g.alpha != 1.0 or g.beta != 1.0 or g.bias is None:
        _fail(f"{node}: unexpected Gemm attributes")
    weight = np.asarray(g.weight, np.float32)
    weight = weight.T if g.trans_b else weight
    out = w.node(node).outputs[0]
    relus = [c for c in w.consumers(out) if c.op_type == "Relu"]
    relu = len(relus) == 1 and len(w.consumers(out)) == 1
    return Linear(name, np.ascontiguousarray(weight), np.asarray(g.bias, np.float32), relu, node,
                  relus[0].outputs[0] if relu else out)


def _conv(w: OnnxWeights, node: str, name: str) -> Conv:
    c = w.conv(node)
    k = c.kernel_shape
    if (c.group != 1 or set(c.dilations) != {1} or len(k) != 2 or k[0] != k[1] or len(set(c.strides)) != 1
            or len(set(c.pads)) != 1 or c.bias is None or c.auto_pad not in ("NOTSET", "")):
        _fail(f"{node}: unexpected Conv attributes {c.kernel_shape} {c.strides} {c.pads} {c.dilations} {c.group}")
    out = w.node(node).outputs[0]
    cons = w.consumers(out)
    act, tensor = None, out
    if len(cons) == 1 and cons[0].op_type in ("HardSwish", "Sigmoid"):
        act, tensor = cons[0].op_type.lower(), cons[0].outputs[0]
    return Conv(name, np.asarray(c.weight, np.float32), np.asarray(c.bias, np.float32), int(c.strides[0]),
                int(c.pads[0]), act, node, tensor)


def _producer_op(w: OnnxWeights, tensor: str) -> str:
    p = w.producer(tensor)
    return p.op_type if p is not None else ""


def _scatter_max_of(w: OnnxWeights, tensor: str) -> str:
    found = [c for c in w.consumers(tensor) if c.op_type == "ScatterElements" and c.attrs.get("reduction") == "max"]
    if len(found) != 1:
        _fail(f"{tensor} does not feed exactly one ScatterElements(max)")
    return found[0].outputs[0]


def load_weights(model_dir: Any, variant: str, *, check_sha256: bool = True) -> FRNetWeights:
    """Parse ``frnet_<variant>.onnx`` of a weights directory (module docstring)."""
    path = variant_files(model_dir, variant)["onnx"]
    digest = file_sha256(path)
    if check_sha256 and digest != ONNX_SHA256[variant]:     # refused before it is parsed
        raise ValueError(f"{path.name}: sha256 {digest} is not the pinned v2.0 file {ONNX_SHA256[variant]}")
    w = OnnxWeights(path)
    if w.input_names != ("points", "coors", "voxel_coors", "inverse_map") or w.output_names != ("pred_probs",):
        _fail(f"inputs {w.input_names} / outputs {w.output_names}")
    taps: Dict[str, str] = {}

    # ---- frustum grids of the six frustum2pixel (ScatterND) sites, in graph order
    grids = []
    for sc in w.nodes("ScatterND"):
        shape_node = w.producer(w.producer(sc.inputs[0]).inputs[0])
        dims = [w.array(t).tolist() for t in shape_node.inputs[:3] if w.has(t)]
        if len(dims) != 3 or dims[0] != [1]:
            _fail(f"{sc.name}: cannot read its (1, H, W, C) image shape")
        grids.append((int(dims[1][0]), int(dims[2][0])))
    if len(grids) != 2 + len(STRIDES):
        _fail(f"{len(grids)} ScatterND sites, expected {2 + len(STRIDES)}")
    H, W = grids[0]
    strides = tuple(H // h for h, _ in grids[2:])
    if grids[1] != (H, W) or strides != STRIDES or any(W // s != g[1] or H // s != g[0]
                                                       for s, g in zip(strides, grids[2:])):
        _fail(f"frustum grids {grids} do not follow strides {STRIDES}")

    # ---- encoder
    bn_pre = w.find_node(f"{VE}/pre_norm/*", "BatchNormalization")
    pre = w.batchnorm(bn_pre.name)
    taps["encoder.input"] = bn_pre.outputs[0]
    ffe = [_point_layer(w, f"{VE}/ffe_layers.{i}/ffe_layers.{i}.0/MatMul", f"encoder.ffe{i}") for i in range(3)]
    for i in range(1, 3):
        if w.node(ffe[i].onnx_node).inputs[0] != ffe[i - 1].onnx_output:
            _fail(f"ffe_layers.{i} does not consume ffe_layers.{i - 1}")
    if w.node(ffe[0].onnx_node).inputs[0] != bn_pre.outputs[0]:
        _fail("ffe_layers.0 does not consume pre_norm")
    ffe3 = _gemm(w, f"{VE}/ffe_layers.3/Gemm", "encoder.ffe3")
    if ffe3.relu or w.node(ffe3.onnx_node).inputs[0] != ffe[2].onnx_output:
        _fail("ffe_layers.3 is not a plain Gemm on ffe_layers.2")
    ffe.append(ffe3)
    for i, lin in enumerate(ffe):
        taps[f"encoder.pf{i}"] = lin.onnx_output
    taps["encoder.voxel_max"] = _scatter_max_of(w, ffe3.onnx_output)
    compression = _gemm(w, f"{VE}/compression_layers/compression_layers.0/Gemm", "encoder.compression")
    if not compression.relu or w.node(compression.onnx_node).inputs[0] != taps["encoder.voxel_max"]:
        _fail("compression_layers is not Gemm + ReLU on the per-voxel max")
    taps["encoder.voxel_feats"] = compression.onnx_output

    # ---- stem and point stem
    stem = tuple(_conv(w, f"{BB}/stem/stem.{i}/Conv", f"backbone.stem{j}") for j, i in enumerate((0, 3, 6)))
    if any(c.activation != "hardswish" or c.stride != 1 for c in stem):
        _fail("stem convs must be stride 1 + HardSwish")
    taps["backbone.frustum"] = w.node(stem[0].onnx_node).inputs[0]
    if _producer_op(w, taps["backbone.frustum"]) != "Transpose":
        _fail("the stem does not read the transposed frustum image")
    taps["backbone.stem"] = stem[-1].onnx_output
    point_stem = _point_layer(w, f"{BB}/point_stem/point_stem.0/MatMul", "backbone.point_stem")
    cat = w.producer(w.node(point_stem.onnx_node).inputs[0])
    if cat.op_type != "Concat" or len(cat.inputs) != 2 or cat.inputs[1] != ffe3.onnx_output:
        _fail("point_stem input is not concat[pixel2point(stem), ffe_layers.3]")
    taps["backbone.point_stem"] = point_stem.onnx_output
    point_stem_segmax = _scatter_max_of(w, point_stem.onnx_output)
    taps["backbone.segmax0"] = point_stem_segmax
    fusion_stem = _conv(w, f"{BB}/fusion_stem/fusion_stem.0/Conv", "backbone.fusion_stem")
    cat = w.producer(w.node(fusion_stem.onnx_node).inputs[0])
    if fusion_stem.activation != "hardswish" or cat.op_type != "Concat" or cat.inputs[1] != stem[-1].onnx_output:
        _fail("fusion_stem is not HardSwish(conv(concat[frustum2pixel(p0), stem]))")
    taps["backbone.fusion_stem"] = fusion_stem.onnx_output

    # ---- stages
    stages: List[Stage] = []
    x_in, p_prev = fusion_stem.onnx_output, point_stem.onnx_output
    for li in range(1, 5):
        conv1s = w.nodes("Conv", f"{BB}/layer{li}/layer{li}.*/conv1/Conv")
        if len(conv1s) != BLOCKS[li - 1]:
            _fail(f"layer{li} has {len(conv1s)} BasicBlocks, expected {BLOCKS[li - 1]}")
        blocks, outs = [], []
        for bi in range(len(conv1s)):
            pre_b = f"{BB}/layer{li}/layer{li}.{bi}"
            c1 = _conv(w, f"{pre_b}/conv1/Conv", f"backbone.layer{li}.block{bi}.conv1")
            c2 = _conv(w, f"{pre_b}/conv2/Conv", f"backbone.layer{li}.block{bi}.conv2")
            ds = None
            if w.nodes("Conv", f"{pre_b}/downsample/*"):
                ds = _conv(w, f"{pre_b}/downsample/downsample.0/Conv", f"backbone.layer{li}.block{bi}.downsample")
            if c1.activation != "hardswish" or c2.activation is not None or (ds is not None and ds.activation):
                _fail(f"{pre_b}: unexpected activations")
            if w.node(c1.onnx_node).inputs[0] != x_in or (ds is not None and w.node(ds.onnx_node).inputs[0] != x_in):
                _fail(f"{pre_b} does not read the previous block output")
            add = _single(w, c2.onnx_output, "Add")
            idt = ds.onnx_output if ds is not None else x_in
            if set(add.inputs) != {c2.onnx_output, idt}:
                _fail(f"{pre_b}: the residual add does not take conv2 + identity")
            hs = _single(w, add.outputs[0], "HardSwish")
            blocks.append((c1, c2, ds))
            outs.append(hs.outputs[0])
            taps[f"backbone.layer{li}.block{bi}"] = hs.outputs[0]
            x_in = hs.outputs[0]
        x_s = outs[-1]
        taps[f"backbone.layer{li}"] = x_s
        pf = _point_layer(w, f"{BB}/point_fusion_layers.{li - 1}/point_fusion_layers.{li - 1}.0/MatMul",
                          f"backbone.point_fusion{li}")
        cat = w.producer(w.node(pf.onnx_node).inputs[0])
        if cat.op_type != "Concat" or len(cat.inputs) != 2 or cat.inputs[1] != p_prev:
            _fail(f"point_fusion_layers.{li - 1} input is not concat[pixel2point(x), p]")
        seg = _scatter_max_of(w, pf.onnx_output)
        taps[f"backbone.point_fusion{li}"] = pf.onnx_output
        taps[f"backbone.segmax{li}"] = seg
        pix = _conv(w, f"{BB}/pixel_fusion_layers.{li - 1}/pixel_fusion_layers.{li - 1}.0/Conv",
                    f"backbone.pixel_fusion{li}")
        cat = w.producer(w.node(pix.onnx_node).inputs[0])
        if pix.activation != "hardswish" or cat.op_type != "Concat" or cat.inputs[1] != x_s:
            _fail(f"pixel_fusion_layers.{li - 1} is not HardSwish(conv(concat[frustum2pixel(p), x]))")
        a0 = _conv(w, f"{BB}/attention_layers.{li - 1}/attention_layers.{li - 1}.0/Conv",
                   f"backbone.attention{li}.conv0")
        a1 = _conv(w, f"{BB}/attention_layers.{li - 1}/attention_layers.{li - 1}.3/Conv",
                   f"backbone.attention{li}.conv1")
        if (a0.activation != "hardswish" or a1.activation != "sigmoid" or w.node(a0.onnx_node).inputs[0] !=
                pix.onnx_output or w.node(a1.onnx_node).inputs[0] != a0.onnx_output):
            _fail(f"attention_layers.{li - 1} is not sigmoid(conv(hswish(conv(pixel_fusion))))")
        mul = _single(w, a1.onnx_output, "Mul")
        if set(mul.inputs) != {pix.onnx_output, a1.onnx_output}:
            _fail(f"stage {li}: the attention product is not pixel_fusion * attention")
        add = _single(w, mul.outputs[0], "Add")
        if set(add.inputs) != {mul.outputs[0], x_s}:
            _fail(f"stage {li}: the output is not pixel_fusion * attention + x")
        taps[f"backbone.pixel_fusion{li}"] = pix.onnx_output
        taps[f"backbone.attention{li}"] = a1.onnx_output
        taps[f"backbone.stage{li}"] = add.outputs[0]
        stride = STRIDES[li - 1]
        stages.append(Stage(li, stride, tuple(blocks), tuple(outs), pf, pix, (a0, a1), seg, add.outputs[0]))
        x_in, p_prev = add.outputs[0], pf.onnx_output
    got = [1] + [st.blocks[0][0].stride for st in stages[1:]]
    if tuple(int(np.prod(got[: i + 1])) for i in range(4)) != STRIDES:
        _fail(f"stage strides {got} do not give {STRIDES}")

    # ---- decoder
    fuse1 = _conv(w, f"{BB}/fuse_layer1/conv/Conv", "decoder.fuse1")
    fuse2 = _conv(w, f"{BB}/fuse_layer2/conv/Conv", "decoder.fuse2")
    cat = w.producer(w.node(fuse1.onnx_node).inputs[0])
    sources = [fusion_stem.onnx_output] + [st.output for st in stages]
    if cat.op_type != "Concat" or len(cat.inputs) != 5:
        _fail("fuse_layer1 does not read a 5-way concat")
    for k, (t, src) in enumerate(zip(cat.inputs, sources)):
        rs = w.producer(t)
        if rs.op_type != "Resize" or rs.inputs[0] != src or rs.attrs.get("mode") != "linear" or \
                rs.attrs.get("coordinate_transformation_mode") != "align_corners":
            _fail(f"decoder input {k} is not an align_corners bilinear Resize of {src}")
        taps[f"decoder.up{k}"] = t
    if fuse1.activation != "hardswish" or fuse2.activation != "hardswish" or \
            w.node(fuse2.onnx_node).inputs[0] != fuse1.onnx_output:
        _fail("fuse_layer1 / fuse_layer2 are not HardSwish convs in sequence")
    taps["decoder.fuse1"] = fuse1.onnx_output
    taps["decoder.fuse2"] = fuse2.onnx_output
    pf1 = _point_layer(w, f"{BB}/point_fuse_layer1/point_fuse_layer1.0/MatMul", "decoder.point_fuse1")
    pf2 = _point_layer(w, f"{BB}/point_fuse_layer2/point_fuse_layer2.0/MatMul", "decoder.point_fuse2")
    cat = w.producer(w.node(pf1.onnx_node).inputs[0])
    if cat.op_type != "Concat" or list(cat.inputs) != [point_stem.onnx_output] + [s.point_fusion.onnx_output
                                                                                 for s in stages]:
        _fail("point_fuse_layer1 input is not concat[p0, p1, p2, p3, p4]")
    if w.node(pf2.onnx_node).inputs[0] != pf1.onnx_output:
        _fail("point_fuse_layer2 does not consume point_fuse_layer1")
    taps["decoder.point_fuse1"] = pf1.onnx_output
    taps["decoder.point_fuse2"] = pf2.onnx_output

    # ---- head: h = pixel2point(fuse2); h = mlp_i(h) + skip_i; logits; softmax
    mlps = [_point_layer(w, f"/model/mlps.{i}/mlps.{i}.0/MatMul", f"head.mlp{i}") for i in range(4)]
    skips = [pf2.onnx_output, ffe[2].onnx_output, ffe[1].onnx_output, ffe[0].onnx_output]
    prev = None
    for i, (lin, skip) in enumerate(zip(mlps, skips)):
        add = _single(w, lin.onnx_output, "Add")
        if set(add.inputs) != {lin.onnx_output, skip}:
            _fail(f"head mlps.{i} skip is not {skip}")
        if prev is not None and w.node(lin.onnx_node).inputs[0] != prev:
            _fail(f"head mlps.{i} does not consume the previous residual sum")
        taps[f"head.mlp{i}"] = add.outputs[0]
        prev = add.outputs[0]
    conv_seg = _gemm(w, "/model/conv_seg/Gemm", "head.logits")
    if conv_seg.relu or w.node(conv_seg.onnx_node).inputs[0] != prev or conv_seg.out_features != NUM_CLASSES:
        _fail("conv_seg is not the final 64 -> 27 Gemm on the last residual sum")
    sm = _single(w, conv_seg.onnx_output, "Softmax")
    if sm.attrs.get("axis", -1) not in (-1, 1) or sm.outputs[0] != "pred_probs":
        _fail("the output is not Softmax(axis=-1) of conv_seg")
    taps["head.logits"] = conv_seg.onnx_output
    taps["pred_probs"] = "pred_probs"

    scale = pre.scale.astype(np.float32)
    shift = pre.shift.astype(np.float32)
    raw = int(sum(int(np.asarray(w.array(t)).size) for t in w.initializer_names()
                  if np.issubdtype(np.asarray(w.array(t)).dtype, np.floating)))
    out = FRNetWeights(variant=variant, frustum_hw=(H, W), strides=strides, pre_norm=pre, pre_norm_scale=scale,
                       pre_norm_shift=shift, ffe=tuple(ffe), compression=compression, stem=stem,
                       point_stem=point_stem, point_stem_segmax=point_stem_segmax, fusion_stem=fusion_stem,
                       stages=tuple(stages), fuse_layer1=fuse1, fuse_layer2=fuse2, point_fuse=(pf1, pf2),
                       head_mlps=tuple(mlps), conv_seg=conv_seg, onnx_taps=taps, onnx_path=str(path),
                       onnx_sha256=digest, raw_parameters=raw)
    folded = sum(c.weight.size + c.bias.size for c in out.convs) + sum(
        lin.weight.size + lin.bias.size for lin in out.linears) + 16
    object.__setattr__(out, "folded_parameters", int(folded))
    missing = [t for t in MODULE_TAPS if t not in taps]
    if missing:
        _fail(f"no ONNX tensor for taps {missing}")
    return out


def check_variant_grid(weights: FRNetWeights, frustum_h: int, frustum_w: int) -> None:
    """The ml_package grid must equal the grid baked into the ONNX (S:frnet:95-98: 1024 OT128 / 256 QT128)."""
    if weights.frustum_hw != (frustum_h, frustum_w):
        raise ValueError(f"ml_package frustum {frustum_h}x{frustum_w} != the ONNX's {weights.frustum_hw}")