code/tt_bevfusion/reference/weights.py

# SPDX-License-Identifier: Apache-2.0
"""The deployed ``bevfusion_lidar.onnx`` of AutowareFoundation/bevfusion@v2.0 read as DATA, the way the TT port reads
it: the vendored ``ttaw.weights.OnnxWeights`` (initializers and Constant tensors addressed by the node that consumes
them; nothing is executed), the module structure checked against the graph topology, every BatchNormalization
folded in float64 with one final rounding (the export keeps all 44 BNs unfolded; the port folds each into its
conv / linear, so this reference does too).

The export (PyTorch 2.2, opset 17 + domain ``autoware``) names its nodes after the mmdet3d modules. Every assumption
the reference makes about the wiring is asserted from the graph (S:bevfusion:205-264):

- sparse encoder (``/pts_middle_encoder``): 21 ``autoware::ImplicitGemm`` (KRSC filters ``[Cout, k0, k1, k2, Cin]``,
  no bias, no fused activation) each followed by ``BatchNormalization`` (eps 1e-3, **folded here**: ``W' = W * s``
  along Cout, ``b' = beta - mean * s``), the basic blocks' residual ``Add`` and ``Relu``; their
  ``GetIndicePairsImplicitGemm`` attributes (subm, ksize, stride, padding, spatial_shape) must be the layer plan of
  ``host/preprocess.py`` (4 sub-manifold levels, 3 stride-2 convs, the 1x1x3 / (1, 1, 2) conv_out); the VFE
  (``ReduceSum / Cast / Div``), the ``coors`` flip and the to-dense ``ScatterND`` -> ``Transpose(0, 4, 3, 1, 2)``
  -> ``Reshape [1, C * Z, X, Y]`` (channel ``c * 2 + z``);
- SECOND (``/pts_backbone/blocks.{0,1}``): 6 + 6 Conv 3x3 (no bias) + BN (eps 1e-3, folded) + Relu, stride 2 first in
  block 1; SECONDFPN (``/pts_neck``): Conv 1x1 + BN + Relu, ConvTranspose 2x2 s2 + BN + Relu (folded), concat;
- head (``/bbox_head``): shared conv 3x3 512 -> 128 (+bias, no activation), heatmap head (3x3 + BN eps 1e-5 (folded)
  + Relu, 3x3 -> 5 + bias), Sigmoid, MaxPool 3x3 s1 p0 on classes 0-3 scattered into a zero map (class 4 passes),
  ``heat * (heat == local_max)``, TopK ``K`` = 500 over the flattened 5 x 32,400, ``class = idx // 32400``, the class
  encoding Conv1d 5 -> 128 of a one-hot, the constant ``bev_pos`` ``(i + 0.5, j + 0.5)`` at ``k = i * 180 + j``, one
  decoder layer (position-embedding MLPs Conv1d 2 -> 128 + BN (folded) + Relu + Conv1d 128 -> 128; 8-head self- and
  cross-attention from ``in_proj_weight`` slices with a ``/ 4`` pre-scale; post-LayerNorms eps 1e-5; FFN
  128 -> 256 -> 128) and six prediction heads (Conv1d 128 -> 64 + BN (folded) + Relu, Conv1d 64 -> C);
- outputs: ``bbox_pred = concat[center + query_pos, height, dim, rot, vel]`` (10, K), ``score = max_c(sigmoid(
  heatmap) * query_heat * onehot(label))`` (K,) and ``label_pred = class`` (K,) int64.

:attr:`BEVFusionWeights.onnx_taps` maps every reference tap name of the dense part to the ONNX tensor that holds the
same value, so the ONNX Runtime golden of the plugin-free dense sub-graph can expose matching intermediates
(``reference/ort.py``). The sparse encoder has no ORT counterpart (``autoware::`` plugins): its oracle is the
research reference (SPEC 6.2).
"""
from __future__ import annotations

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

import numpy as np

from ..ttaw.weights import OnnxWeights, file_sha256, fold_bn, fold_bn_conv, fold_bn_conv_transpose
from .config import BEV_SHAPE, NUM_CLASSES, ONNX_FILE, SPARSE_LEVELS, BEVFusionConfig

__all__ = ["SparseLayer", "Layer", "AttentionParams", "DecoderParams", "HeadParams", "BEVFusionWeights",
           "load_weights", "PREDICTION_HEADS", "ENCODER_PLAN", "PME", "BB", "NK", "HB", "DEC", "DENSE_INPUT"]

PME, BB, NK, HB = "/pts_middle_encoder", "/pts_backbone", "/pts_neck", "/bbox_head"
DEC = f"{HB}/decoder.0"
DENSE_INPUT = f"{PME}/Reshape_3_output_0"     # the to-dense BEV [1, 256, 180, 180]: input of the dense sub-graph
# prediction heads in ONNX output order, channels
PREDICTION_HEADS = (("center", 2), ("height", 1), ("dim", 3), ("rot", 2), ("vel", 2), ("heatmap", NUM_CLASSES))

# The sparse encoder, layer by layer: (tap name, ONNX module path, neighbour map of host/preprocess.py, Cin, Cout,
# GetIndicePairs attributes (subm, ksize, stride, padding, spatial_shape), residual source tap or None).
def _encoder_plan() -> List[Tuple[str, str, str, int, int, tuple, Optional[str]]]:
    p: List[Tuple[str, str, str, int, int, tuple, Optional[str]]] = []
    shapes = [SPARSE_LEVELS[k] for k in ("l1", "l2", "l3", "l4")]
    p.append(("sparse.conv_input", f"{PME}/conv_input/conv_input.0", "subm1", 4, 16,
              (1, (3, 3, 3), (1, 1, 1), (1, 1, 1), shapes[0]), None))
    chans = (16, 32, 64, 128)
    prev = "sparse.conv_input"
    for s in range(1, 5):
        c = chans[s - 1]
        for b in range(2):
            pre = f"{PME}/encoder_layer{s}/encoder_layer{s}.{b}"
            attrs = (1, (3, 3, 3), (1, 1, 1), (1, 1, 1), shapes[s - 1])
            p.append((f"sparse.layer{s}.{b}.conv1", f"{pre}/conv1", f"subm{s}", c, c, attrs, None))
            p.append((f"sparse.layer{s}.{b}", f"{pre}/conv2", f"subm{s}", c, c, attrs, prev))
            prev = f"sparse.layer{s}.{b}"
        if s < 4:
            pad = (1, 1, 0) if s == 3 else (1, 1, 1)
            p.append((f"sparse.down{s}", f"{PME}/encoder_layer{s}/encoder_layer{s}.2/encoder_layer{s}.2.0",
                      f"down{s}", c, chans[s], (0, (3, 3, 3), (2, 2, 2), pad, shapes[s - 1]), None))
            prev = f"sparse.down{s}"
    p.append(("sparse.conv_out", f"{PME}/conv_out/conv_out.0", "conv_out", 128, 128,
              (0, (1, 1, 3), (1, 1, 2), (0, 0, 0), shapes[3]), None))
    return p


ENCODER_PLAN = _encoder_plan()


@dataclass(frozen=True)
class SparseLayer:
    """One sparse conv, BN folded: ``y = relu(gather_conv(x, weight) + bias (+ residual))``."""

    name: str                 # tap name ("sparse.layer2.1", "sparse.down3", ...)
    map: str                  # neighbour map of the host plan ("subm2", "down3", "conv_out", ...)
    weight: np.ndarray        # float32 KRSC [Cout, k0, k1, k2, Cin], BN folded
    bias: np.ndarray          # float32 [Cout] (the folded BN shift)
    residual: Optional[str]   # tap name of the identity added before the Relu (block conv2), or None
    onnx_node: str            # the ImplicitGemm node
    bn_node: str

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

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

    @property
    def kernel_volume(self) -> int:
        return int(np.prod(self.weight.shape[1:4]))

    @property
    def num_params(self) -> int:
        return int(self.weight.size + self.bias.size)


@dataclass(frozen=True)
class Layer:
    """One dense conv layer, BN folded: ``y = relu?(conv(x, weight) + bias)`` (2-D, or 1-D for the head's Conv1d)."""

    name: str                 # tap name ("backbone.block0.conv0", "neck.deblock1", "heads.center.hidden", ...)
    kind: str                 # "conv" | "conv_transpose" | "conv1d"
    weight: np.ndarray        # float32: conv [Cout, Cin, k, k]; conv_transpose [Cin, Cout, k, k]; conv1d [Cout, Cin, 1]
    bias: np.ndarray          # float32 [Cout]
    stride: int
    padding: int
    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 kernel(self) -> int:
        return int(self.weight.shape[-1])

    @property
    def in_channels(self) -> int:
        return int(self.weight.shape[0] if self.kind == "conv_transpose" else self.weight.shape[1])

    @property
    def out_channels(self) -> int:
        return int(self.weight.shape[1] if self.kind == "conv_transpose" else self.weight.shape[0])

    @property
    def num_params(self) -> int:
        return int(self.weight.size + self.bias.size)


@dataclass(frozen=True)
class AttentionParams:
    """``nn.MultiheadAttention`` as exported: ``q = x_q @ wq + bq`` (``[in, out]``; the rows of ``in_proj_weight``
    sliced and transposed by the graph), likewise k / v, per head ``softmax((q / div) @ k^T) @ v``, ``out =
    concat_heads @ wo + bo`` (``wo`` stored ``[in, out]`` here; the ONNX Gemm keeps ``[out, in]`` with transB=1)."""

    wq: np.ndarray
    bq: np.ndarray
    wk: np.ndarray
    bk: np.ndarray
    wv: np.ndarray
    bv: np.ndarray
    wo: np.ndarray
    bo: np.ndarray
    num_heads: int
    head_dim: int
    q_divisor: float          # the ONNX pre-scale Div constant (sqrt(head_dim) = 4.0)

    @property
    def scale(self) -> float:
        return 1.0 / self.q_divisor


@dataclass(frozen=True)
class DecoderParams:
    """One TransFusion decoder layer (post-norm): self-attention on ``query + query_pos_embed`` (Q = K = V),
    cross-attention ``query = x + query_pos_embed`` over ``key = value = lidar_feat + key_pos_embed``, FFN."""

    self_posembed: Tuple[Layer, Layer]       # Conv1d 2 -> 128 + BN (folded) + Relu, Conv1d 128 -> 128
    cross_posembed: Tuple[Layer, Layer]
    self_attn: AttentionParams
    cross_attn: AttentionParams
    norms: Tuple[Tuple[np.ndarray, np.ndarray, float], ...]    # (gamma, beta, eps) x 3
    ffn_w1: np.ndarray        # (128, 256) [in, out]
    ffn_b1: np.ndarray
    ffn_w2: np.ndarray        # (256, 128)
    ffn_b2: np.ndarray


@dataclass(frozen=True)
class HeadParams:
    shared: Layer                                # 3x3 512 -> 128 + bias, no activation (lidar_feat)
    heatmap_hidden: Layer                        # 3x3 128 -> 128 + BN (folded) + Relu
    heatmap_out: Layer                           # 3x3 128 -> 5 (logits)
    nms_kernel: int                              # 3 (MaxPool kernel; stride 1, pad 0, zero border on pooled classes)
    pooled_classes: Tuple[int, ...]              # (0, 1, 2, 3): class 4 (PEDESTRIAN) is not pooled
    num_proposals: int                           # TopK K = 500
    class_encoding_w: np.ndarray                 # (128, 5) = Conv1d weight[:, :, 0]
    class_encoding_b: np.ndarray                 # (128,)
    bev_pos: np.ndarray                          # (32400, 2) float32 constant: (i + 0.5, j + 0.5), k = i * 180 + j
    decoder: DecoderParams
    prediction_heads: Tuple[Tuple[str, Layer, Layer], ...]   # (name, Conv1d 128 -> 64 + BN + Relu, Conv1d 64 -> C)

    @property
    def class_table(self) -> np.ndarray:
        """``(5, 128)`` float32: row c = ``class_encoding(one_hot(c))`` = ``W[:, c] + b``."""
        return (self.class_encoding_w.T.astype(np.float32) + self.class_encoding_b.astype(np.float32)).astype(
            np.float32)


@dataclass(frozen=True)
class BEVFusionWeights:
    encoder: Tuple[SparseLayer, ...]             # 21 sparse convs in forward order (ENCODER_PLAN)
    blocks: Tuple[Tuple[Layer, ...], ...]        # SECOND: 2 blocks of 6 3x3 convs
    deblocks: Tuple[Layer, ...]                  # SECONDFPN: conv 1x1, convT 2x2 s2
    head: HeadParams
    onnx_taps: Dict[str, str]                    # reference tap name -> ONNX tensor name (dense part)
    sha256: str
    path: str

    def dense_layers(self) -> List[Layer]:
        """Every 2-D conv layer (backbone, neck, head) in forward order."""
        h = self.head
        return ([layer for block in self.blocks for layer in block] + list(self.deblocks)
                + [h.shared, h.heatmap_hidden, h.heatmap_out])

    def num_params(self) -> int:
        """Learned parameters after the BN folds (the ONNX has 8,308,692 float32 weights incl. the BN statistics)."""
        h = self.head
        d = h.decoder
        n = sum(x.num_params for x in self.encoder) + sum(x.num_params for x in self.dense_layers())
        n += h.class_encoding_w.size + h.class_encoding_b.size
        for pe in (d.self_posembed, d.cross_posembed):
            n += sum(x.num_params for x in pe)
        for a in (d.self_attn, d.cross_attn):
            n += sum(x.size for x in (a.wq, a.bq, a.wk, a.bk, a.wv, a.bv, a.wo, a.bo))
        n += sum(g.size + b.size for g, b, _ in d.norms)
        n += sum(x.size for x in (d.ffn_w1, d.ffn_b1, d.ffn_w2, d.ffn_b2))
        for _, hid, fin in h.prediction_heads:
            n += hid.num_params + fin.num_params
        return int(n)


# ------------------------------------------------------------------------------------------------- helpers

def _f32(a: Any) -> np.ndarray:
    return np.asarray(a, dtype=np.float32)


def _single(found: Sequence[Any], what: str) -> Any:
    if len(found) != 1:
        raise ValueError(f"expected exactly one {what}, found {len(found)}")
    return found[0]


def _expect(cond: bool, msg: str) -> None:
    if not cond:
        raise ValueError(f"bevfusion_lidar.onnx structure: {msg}")


def _consumer(w: OnnxWeights, tensor: str, op_type: str) -> Any:
    cons = [n for n in w.consumers(tensor) if n.op_type == op_type]
    return _single(cons, f"{op_type} consumer of {tensor!r}")


def _scalar(w: OnnxWeights, tensor: str) -> float:
    a = np.asarray(w.array(tensor)).reshape(-1)
    _expect(a.size == 1, f"{tensor} is not a scalar constant")
    return float(a[0])


def _fold_relu_chain(w: OnnxWeights, node_name: str) -> Tuple[Any, Any, Any]:
    """``node -> BatchNormalization -> Relu`` (the single consumers): returns (node, bn, relu)."""
    node = w.node(node_name)
    bn = _consumer(w, node.outputs[0], "BatchNormalization")
    relu = _consumer(w, bn.outputs[0], "Relu")
    return node, bn, relu


def _conv_bn_relu(w: OnnxWeights, name: str, node_name: str, *, kind: str = "conv") -> Layer:
    """Conv / ConvTranspose (no bias) + BatchNormalization (folded) + Relu."""
    node, bn, relu = _fold_relu_chain(w, node_name)
    c = w.conv(node_name)
    _expect(c.group == 1 and all(d == 1 for d in c.dilations) and len(set(c.strides)) == 1
            and len(set(c.pads)) <= 1 and len(set(c.kernel_shape)) == 1 and c.auto_pad == "NOTSET",
            f"{node_name}: unexpected attributes {c}")
    b = None if c.bias is None else np.asarray(c.bias, np.float64)
    if kind == "conv_transpose":
        _expect(tuple(c.kernel_shape) == tuple(c.strides) and not any(c.pads), f"{node_name}: ConvTranspose k != s")
        wf, bf = fold_bn_conv_transpose(np.asarray(c.weight, np.float64), b, w.batchnorm(bn.name))
    else:
        wf, bf = fold_bn_conv(np.asarray(c.weight, np.float64), b, w.batchnorm(bn.name))
    pad = int(c.pads[0]) if c.pads else 0
    return Layer(name, kind, _f32(wf), _f32(bf), int(c.strides[0]), pad, True, node_name, relu.outputs[0],
                 bn_node=bn.name)


def _conv_plain(w: OnnxWeights, name: str, node_name: str, *, kind: str = "conv") -> Layer:
    """Conv with bias, no BN, no activation."""
    c = w.conv(node_name)
    _expect(c.group == 1 and all(d == 1 for d in c.dilations) and len(set(c.strides)) == 1
            and len(set(c.pads)) <= 1, f"{node_name}: unexpected attributes {c}")
    bias = np.zeros(c.weight.shape[0], np.float32) if c.bias is None else _f32(c.bias)
    pad = int(c.pads[0]) if c.pads else 0
    return Layer(name, kind, _f32(c.weight), bias, int(c.strides[0]), pad, False, node_name,
                 w.node(node_name).outputs[0])


# ----------------------------------------------------------------------------------------- sparse encoder

def _load_encoder(w: OnnxWeights) -> Tuple[Tuple[SparseLayer, ...], Dict[str, str]]:
    # VFE and coordinate flip feeding the first conv: features = ReduceSum(voxels, 1) / float(num_points)
    rs = w.node("/ReduceSum")
    div = w.node("/Div")
    _expect(rs.inputs[0] == "voxels" and int(rs.attrs.get("keepdims", 1)) == 0
            and list(np.asarray(w.array(rs.inputs[1])).reshape(-1)) == [1] and div.inputs[0] == rs.outputs[0],
            "VFE is not ReduceSum(voxels, axis=1) / num_points")
    sl = w.node("/Slice")
    _expect(sl.inputs[0] == "coors" and int(np.asarray(w.array(sl.inputs[3])).reshape(-1)[0]) == -1,
            "coors are not flipped (z, y, x) -> (x, y, z)")
    layers: List[SparseLayer] = []
    taps: Dict[str, str] = {}
    expected_input = div.outputs[0]
    relu_out: Dict[str, str] = {}
    for name, path, map_name, cin, cout, attrs, residual in ENCODER_PLAN:
        gemm = _single(w.nodes("ImplicitGemm", f"{path}/ImplicitGemm"), f"ImplicitGemm of {path}")
        pairs = w.node(f"{path}/GetIndicePairsImplicitGemm")
        a = pairs.attrs
        got = (int(a["subm"]), tuple(int(v) for v in a["ksize"]), tuple(int(v) for v in a["stride"]),
               tuple(int(v) for v in a["padding"]), tuple(int(v) for v in a["spatial_shape"]))
        _expect(got == attrs, f"{pairs.name}: attributes {got} != the plan {attrs}")
        _expect(tuple(int(v) for v in a["dilation"]) == (1, 1, 1) and int(a.get("transpose", 0)) == 0
                and int(a.get("batch_size", 1)) == 1, f"{pairs.name}: dilation / transpose / batch")
        _expect(int(gemm.attrs.get("is_subm", -1)) == attrs[0] and float(gemm.attrs.get("output_scale", 1.0)) == 1.0
                and float(gemm.attrs.get("output_add_scale", 0.0)) == 0.0 and len(gemm.inputs) == 5,
                f"{gemm.name}: fused bias / activation / scale")
        _expect(gemm.inputs[0] == expected_input, f"{gemm.name} does not read the previous layer ({expected_input})")
        wt = np.asarray(w.array(gemm.inputs[1]), np.float64)
        _expect(wt.shape == (cout, *attrs[1], cin), f"{gemm.name}: filter {wt.shape} != KRSC {(cout, *attrs[1], cin)}")
        bn = _consumer(w, gemm.outputs[0], "BatchNormalization")
        bnp = w.batchnorm(bn.name)
        _expect(abs(bnp.eps - 1e-3) < 1e-9, f"{bn.name}: eps {bnp.eps}")
        wf, bf = fold_bn(wt, None, bnp, axis=0)
        nxt = bn.outputs[0]
        if residual is not None:
            add = _consumer(w, nxt, "Add")
            _expect(set(add.inputs) == {nxt, relu_out[residual]}, f"{add.name}: residual is not {residual}")
            nxt = add.outputs[0]
        relu = _consumer(w, nxt, "Relu")
        layers.append(SparseLayer(name, map_name, _f32(wf), _f32(bf), residual, gemm.name, bn.name))
        relu_out[name] = relu.outputs[0]
        taps[name] = relu.outputs[0]
        expected_input = relu.outputs[0]
    # to-dense: ScatterND of the conv_out features into [1, X, Y, Z, C], Transpose(0, 4, 3, 1, 2), Reshape
    scat = _single(w.nodes("ScatterND", f"{PME}/*"), "to-dense ScatterND")
    _expect(w.producer(scat.inputs[2]) is not None, "to-dense updates")
    t1 = w.node(f"{PME}/Transpose_1")
    _expect(t1.inputs[0] == scat.outputs[0] and list(t1.attrs["perm"]) == [0, 4, 3, 1, 2],
            "to-dense transpose is not (0, 4, 3, 1, 2): channel order c * Z + z")
    rs3 = w.node(f"{PME}/Reshape_3")
    _expect(rs3.inputs[0] == t1.outputs[0] and rs3.outputs[0] == DENSE_INPUT, "to-dense reshape")
    _expect(len(layers) == 21, f"{len(layers)} sparse convs, expected 21")
    taps["sparse.bev"] = DENSE_INPUT
    return tuple(layers), taps


# ------------------------------------------------------------------------------------- SECOND / SECONDFPN

def _load_backbone_neck(w: OnnxWeights) -> Tuple[tuple, tuple, Dict[str, str]]:
    taps: Dict[str, str] = {}
    blocks, prev = [], DENSE_INPUT
    for b in range(2):
        convs = w.nodes("Conv", f"{BB}/blocks.{b}/*")
        layers = []
        for j, conv in enumerate(convs):
            _expect(conv.inputs[0] == prev, f"{conv.name} does not consume the previous layer output")
            layer = _conv_bn_relu(w, f"backbone.block{b}.conv{j}", conv.name)
            stride = 2 if (b == 1 and j == 0) else 1
            _expect(layer.kernel == 3 and layer.padding == 1 and layer.stride == stride
                    and not np.asarray(w.conv(conv.name).bias if w.conv(conv.name).bias is not None else 0).any(),
                    f"{conv.name}: SECOND conv must be 3x3 pad 1 (stride {stride}), no bias")
            layers.append(layer)
            taps[layer.name] = layer.onnx_output
            prev = layer.onnx_output
        blocks.append(tuple(layers))
        taps[f"backbone.block{b}"] = prev
    _expect([len(b) for b in blocks] == [6, 6], f"blocks have {[len(b) for b in blocks]} convs, expected 6 / 6")
    concat = w.node(f"{NK}/Concat")
    _expect(int(concat.attrs.get("axis", -1)) == 1 and len(concat.inputs) == 2, "neck concat")
    d0 = _conv_bn_relu(w, "neck.deblock0", f"{NK}/deblocks.0/deblocks.0.0/Conv")
    d1 = _conv_bn_relu(w, "neck.deblock1", f"{NK}/deblocks.1/deblocks.1.0/ConvTranspose", kind="conv_transpose")
    _expect(list(concat.inputs) == [d0.onnx_output, d1.onnx_output], "neck concat inputs")
    _expect(w.node(d0.onnx_node).inputs[0] == taps["backbone.block0"]
            and w.node(d1.onnx_node).inputs[0] == taps["backbone.block1"], "deblock inputs")
    _expect((d0.kind, d0.kernel, d0.stride, d1.kind, d1.kernel, d1.stride) == ("conv", 1, 1, "conv_transpose", 2, 2),
            "deblocks (expected conv 1x1, convT 2x2 s2)")
    taps.update({"neck.deblock0": d0.onnx_output, "neck.deblock1": d1.onnx_output, "neck": concat.outputs[0]})
    return tuple(blocks), (d0, d1), taps


# ---------------------------------------------------------------------------------------------------- head

def _attention(w: OnnxWeights, prefix: str) -> AttentionParams:
    a = f"{prefix}/attn"
    mats = []
    for i, suffix in enumerate(("", "_1", "_2")):
        mm = w.node(f"{a}/MatMul{suffix}")
        tr = w.producer(mm.inputs[1])
        sl = w.producer(tr.inputs[0])
        _expect(tr.op_type == "Transpose" and sl.op_type == "Slice", f"{mm.name}: weight is not Transpose(Slice)")
        mats.append(sl.inputs[0])
    _expect(len(set(mats)) == 1, f"{a}: Q / K / V do not slice one in_proj_weight")
    W = np.asarray(w.array(mats[0]), np.float32)                       # [3 * E, E]
    E = W.shape[1]
    _expect(W.shape == (3 * E, E), f"{a}: in_proj_weight {W.shape}")
    bias_name = w.producer(w.node(f"{a}/Add_2").inputs[0]).inputs[0]
    B = np.asarray(w.array(bias_name), np.float32)
    _expect(B.shape == (3 * E,), f"{a}: in_proj_bias {B.shape}")
    for n, mm in (("Add_2", "MatMul"), ("Add_3", "MatMul_1"), ("Add_4", "MatMul_2")):
        _expect(w.node(f"{a}/{mm}").outputs[0] in w.node(f"{a}/{n}").inputs, f"{a}/{n} does not add the bias of {mm}")
    g = w.gemm(f"{a}/Gemm")
    _expect(g.trans_b and not g.trans_a and g.alpha == 1.0 and g.beta == 1.0, f"{a}/Gemm attributes")
    heads = int(round(_scalar(w, w.node(f"{a}/Div").inputs[1])))
    hd = E // heads
    div = _scalar(w, w.node(f"{a}/Div_3").inputs[1])
    _expect(heads * hd == E and abs(div - np.sqrt(hd)) < 1e-6, f"{a}: {heads} heads, divisor {div}")
    return AttentionParams(W[:E].T.copy(), B[:E].copy(), W[E:2 * E].T.copy(), B[E:2 * E].copy(),
                           W[2 * E:].T.copy(), B[2 * E:].copy(), _f32(g.weight).T.copy(), _f32(g.bias), heads, hd,
                           div)


def _posembed(w: OnnxWeights, which: str) -> Tuple[Layer, Layer]:
    p = f"{DEC}/{which}/position_embedding_head"
    node, bn, relu = _fold_relu_chain(w, f"{p}/position_embedding_head.0/Conv")
    c = w.conv(node.name)
    wf, bf = fold_bn_conv(np.asarray(c.weight, np.float64), np.asarray(c.bias, np.float64), w.batchnorm(bn.name))
    l0 = Layer(f"decoder.{which}.0", "conv1d", _f32(wf), _f32(bf), 1, 0, True, node.name, relu.outputs[0],
               bn_node=bn.name)
    l1 = _conv_plain(w, f"decoder.{which}.1", f"{p}/position_embedding_head.3/Conv", kind="conv1d")
    _expect(w.node(l1.onnx_node).inputs[0] == l0.onnx_output and l0.weight.shape == (128, 2, 1)
            and l1.weight.shape == (128, 128, 1), f"{which} MLP shape / wiring")
    return l0, l1


def _load_head(w: OnnxWeights, cfg: Optional[BEVFusionConfig], neck_out: str) -> Tuple[HeadParams, Dict[str, str]]:
    taps: Dict[str, str] = {}
    shared = _conv_plain(w, "head.shared", f"{HB}/shared_conv/Conv")
    _expect(w.node(shared.onnx_node).inputs[0] == neck_out and (shared.in_channels, shared.out_channels,
                                                                 shared.kernel) == (512, 128, 3), "shared conv")
    cast1 = w.node(f"{HB}/Cast_1")
    _expect(cast1.inputs[0] == shared.onnx_output, "heatmap head input is not the shared conv")
    hh = _conv_bn_relu(w, "head.heatmap.hidden", f"{HB}/heatmap_head/heatmap_head.0/conv/Conv")
    _expect(abs(w.batchnorm(hh.bn_node).eps - 1e-5) < 1e-9, "heatmap head BN eps != 1e-5")
    ho = _conv_plain(w, "head.heatmap", f"{HB}/heatmap_head/heatmap_head.1/Conv")
    _expect(w.node(hh.onnx_node).inputs[0] == cast1.outputs[0] and w.node(ho.onnx_node).inputs[0] == hh.onnx_output
            and ho.out_channels == NUM_CLASSES and hh.kernel == ho.kernel == 3, "heatmap head")
    sig = w.node(f"{HB}/Sigmoid")
    _expect(sig.inputs[0] == ho.onnx_output, "sigmoid of the heatmap")
    # local max: MaxPool 3x3 s1 p0 of the pooled classes, scattered into a zero map; the other classes pass
    g2 = w.node(f"{HB}/Gather_2")
    pooled = tuple(int(v) for v in np.asarray(w.array(g2.inputs[1])).reshape(-1))
    pool = w.node(f"{HB}/MaxPool")
    k = list(pool.attrs.get("kernel_shape", []))
    _expect(g2.inputs[0] == sig.outputs[0] and pool.inputs[0] == g2.outputs[0] and len(set(k)) == 1
            and list(pool.attrs.get("strides", [1, 1])) == [1, 1] and not any(pool.attrs.get("pads", [0, 0, 0, 0])),
            f"MaxPool wiring / attributes {pool.attrs}")
    g6 = w.node(f"{HB}/Gather_6")
    passing = tuple(int(v) for v in np.asarray(w.array(g6.inputs[1])).reshape(-1))
    _expect(pooled == (0, 1, 2, 3) and passing == (4,) and g6.inputs[0] == sig.outputs[0],
            f"pooled classes {pooled}, passing {passing}")
    eq = w.node(f"{HB}/Equal_6")
    mul6 = w.node(f"{HB}/Mul_6")
    scat1 = w.node(f"{HB}/ScatterND_1")
    _expect(list(eq.inputs) == [sig.outputs[0], scat1.outputs[0]] and sig.outputs[0] in mul6.inputs,
            "heat * (heat == local_max) wiring")
    topk = w.node(f"{HB}/TopK")
    K = int(np.asarray(w.array(topk.inputs[1])).reshape(-1)[0])
    _expect(int(topk.attrs.get("largest", 1)) == 1 and int(topk.attrs.get("sorted", 1)) == 1, "TopK attributes")
    enc = w.conv(f"{HB}/class_encoding/Conv")
    _expect(enc.weight.shape == (128, NUM_CLASSES, 1), f"class encoding {enc.weight.shape}")
    onehot = w.node(f"{HB}/OneHot")
    _expect(int(np.asarray(w.array(onehot.inputs[1])).reshape(-1)[0]) == NUM_CLASSES, "OneHot depth")
    bev_pos = _f32(w.array(f"{HB}/Constant_2_output_0")).reshape(-1, 2)
    H, W_ = BEV_SHAPE
    kk = np.arange(H * W_)
    _expect(bev_pos.shape == (H * W_, 2)
            and np.array_equal(bev_pos[:, 0], (kk // W_).astype(np.float32) + np.float32(0.5))
            and np.array_equal(bev_pos[:, 1], (kk % W_).astype(np.float32) + np.float32(0.5)),
            "bev_pos is not (i + 0.5, j + 0.5) with k = i * 180 + j")
    if cfg is not None:
        _expect(K == cfg.num_proposals, f"ONNX TopK K={K} != ml_package num_proposals {cfg.num_proposals}")
    # decoder
    norms = []
    for i in range(3):
        n = f"{DEC}/norms.{i}"
        eps = _scalar(w, w.node(f"{n}/Add").inputs[1])
        norms.append((_f32(w.param(f"{n}/Mul", 1)), _f32(w.param(f"{n}/Add_1", 1)), eps))
    ffn0 = w.node(f"{DEC}/ffn/layers/layers.0/layers.0.0/MatMul")
    ffn1 = w.node(f"{DEC}/ffn/layers/layers.1/MatMul")
    dec = DecoderParams(
        self_posembed=_posembed(w, "self_posembed"), cross_posembed=_posembed(w, "cross_posembed"),
        self_attn=_attention(w, f"{DEC}/self_attn"), cross_attn=_attention(w, f"{DEC}/cross_attn"),
        norms=tuple(norms),
        ffn_w1=_f32(w.array(w.producer(ffn0.inputs[1]).inputs[0])).T.copy(),
        ffn_b1=_f32(w.param(f"{DEC}/ffn/layers/layers.0/layers.0.0/Add", 0)),
        ffn_w2=_f32(w.array(w.producer(ffn1.inputs[1]).inputs[0])).T.copy(),
        ffn_b2=_f32(w.param(f"{DEC}/ffn/layers/layers.1/Add", 0)))
    _expect(all(abs(eps - 1e-5) < 1e-9 for _, _, eps in dec.norms), "LayerNorm eps != 1e-5")
    _expect(dec.ffn_w1.shape == (128, 256) and dec.ffn_w2.shape == (256, 128), "FFN shapes")
    _expect(w.node(f"{DEC}/ffn/Add").inputs[0] == w.node(f"{DEC}/norms.1/Add_1").outputs[0], "FFN residual")
    # wiring of the decoder inputs: Q = K = V = query + pos in self-attention, cross key = lidar_feat + key pos
    q_in = {w.node(f"{DEC}/self_attn/Add").inputs, w.node(f"{DEC}/self_attn/Add_1").inputs, w.node(f"{DEC}/Add").inputs}
    _expect(len({frozenset(x) for x in q_in}) == 1, "self-attention Q / K / V are not all query + query_pos_embed")
    kv = w.node(f"{DEC}/Add_1")
    kpe = w.node(f"{DEC}/Transpose_1").outputs[0]
    _expect(kpe in kv.inputs and w.node(f"{DEC}/Transpose_3").outputs[0] in kv.inputs,
            "cross-attention key / value are not lidar_feat + key_pos_embed")
    heads = []
    for name, ch in PREDICTION_HEADS:
        p = f"{HB}/prediction_heads.0/{name}"
        node, bn, relu = _fold_relu_chain(w, f"{p}/{name}.0/conv/Conv")
        c = w.conv(node.name)
        b = None if c.bias is None else np.asarray(c.bias, np.float64)
        wf, bf = fold_bn_conv(np.asarray(c.weight, np.float64), b, w.batchnorm(bn.name))
        hid = Layer(f"heads.{name}.hidden", "conv1d", _f32(wf), _f32(bf), 1, 0, True, node.name, relu.outputs[0],
                    bn_node=bn.name)
        fin = _conv_plain(w, f"heads.{name}", f"{p}/{name}.1/Conv", kind="conv1d")
        _expect(hid.weight.shape == (64, 128, 1) and fin.weight.shape == (ch, 64, 1)
                and w.node(fin.onnx_node).inputs[0] == hid.onnx_output, f"prediction head {name}")
        heads.append((name, hid, fin))
        taps[fin.name] = fin.onnx_output
    head = HeadParams(shared=shared, heatmap_hidden=hh, heatmap_out=ho, nms_kernel=int(k[0]), pooled_classes=pooled,
                      num_proposals=K, class_encoding_w=_f32(enc.weight[:, :, 0]), class_encoding_b=_f32(enc.bias),
                      bev_pos=bev_pos, decoder=dec, prediction_heads=tuple(heads))
    taps.update({
        "head.shared": shared.onnx_output, "head.heatmap.hidden": hh.onnx_output, "head.heatmap": ho.onnx_output,
        "head.heat": sig.outputs[0], "head.heat_nms": mul6.outputs[0], "head.topk_values": topk.outputs[0],
        "head.topk_indices": topk.outputs[1], "query.feat": w.node(f"{HB}/Add_4").outputs[0],
        "query.pos": w.node(f"{HB}/GatherElements_1").outputs[0],
        "query.heat": w.node(f"{HB}/GatherElements_2").outputs[0], "decoder.key_pos_embed": kpe,
        "decoder.query_pos_embed": w.node(f"{DEC}/Transpose").outputs[0],
        "decoder.self_attn": w.node(f"{DEC}/norms.0/Add_1").outputs[0],
        "decoder.cross_attn": w.node(f"{DEC}/norms.1/Add_1").outputs[0],
        "decoder.out": w.node(f"{DEC}/norms.2/Add_1").outputs[0],
    })
    _check_outputs(w, heads)
    return head, taps


def _check_outputs(w: OnnxWeights, heads: Sequence[Tuple[str, Layer, Layer]]) -> None:
    """``bbox_pred = concat[center + query_pos, height, dim, rot, vel]`` (batch 0), ``score = ReduceMax_c(sigmoid(
    heatmap) * query_heat * onehot(label))`` (batch 0), ``label_pred = class`` (batch 0)."""
    out = {name: fin.onnx_output for name, _, fin in heads}

    def through(t: str) -> str:                     # follow single-input Concats and the batch-0 Gather forward
        while True:
            cons = [n for n in w.consumers(t) if (n.op_type == "Concat" and len(n.inputs) == 1)
                    or (n.op_type == "Gather" and n.inputs[1] == "/Constant_output_0")]
            if not cons:
                return t
            t = cons[0].outputs[0]

    center_add = w.node(f"{HB}/Add_5")
    qpos_t = w.node(f"{HB}/Transpose_2")
    _expect(list(center_add.inputs) == [out["center"], qpos_t.outputs[0]]
            and qpos_t.inputs[0] == w.node(f"{HB}/GatherElements_1").outputs[0], "center is not offset by query_pos")
    bbox = w.producer("bbox_pred")
    _expect(bbox.op_type == "Concat" and list(bbox.inputs) == [through(center_add.outputs[0]), through(out["height"]),
                                                                through(out["dim"]), through(out["rot"]),
                                                                through(out["vel"])], f"bbox_pred = {bbox.inputs}")
    sig = w.node("/Sigmoid")
    mul = w.node("/Mul")
    _expect(w.producer(sig.inputs[0]).inputs[0] == out["heatmap"]
            and list(mul.inputs) == [sig.outputs[0], w.node(f"{HB}/GatherElements_2").outputs[0]],
            "score is not sigmoid(heatmap) * query_heat")
    rmax = w.producer("score")
    _expect(rmax.op_type == "ReduceMax" and list(rmax.attrs.get("axes", [])) == [0], "score ReduceMax over classes")
    lab = w.producer("label_pred")
    _expect(lab.op_type == "Gather" and lab.inputs[0] == w.node(f"{HB}/Cast_8").outputs[0], "label_pred = class")
    g2 = w.node(f"{HB}/GatherElements_2")
    _expect(w.producer(g2.inputs[0]).inputs[0] == w.node(f"{HB}/Mul_6").outputs[0],
            "query_heat is not gathered from the local-max-masked heatmap (all classes)")


# ---------------------------------------------------------------------------------------------------- load

def load_weights(weights_dir: Any, config: Optional[BEVFusionConfig] = None) -> BEVFusionWeights:
    """``<weights_dir>/bevfusion_lidar.onnx`` (or the ``.onnx`` path itself) -> :class:`BEVFusionWeights`, checked
    against ``config`` (K) when given."""
    p = Path(weights_dir)
    path = p if p.suffix == ".onnx" else p / ONNX_FILE
    w = OnnxWeights(path)
    _expect(tuple(w.input_names) == ("voxels", "coors", "num_points_per_voxel"), f"inputs {w.input_names}")
    _expect(tuple(w.output_names) == ("bbox_pred", "score", "label_pred"), f"outputs {w.output_names}")
    encoder, taps = _load_encoder(w)
    blocks, deblocks, t2 = _load_backbone_neck(w)
    taps.update(t2)
    head, t3 = _load_head(w, config, taps["neck"])
    taps.update(t3)
    taps.update({"bbox_pred": "bbox_pred", "score": "score", "label_pred": "label_pred"})
    return BEVFusionWeights(encoder=encoder, blocks=blocks, deblocks=deblocks, head=head, onnx_taps=taps,
                            sha256=file_sha256(path), path=str(path))