code/tt_bevformer/reference/model.py

# SPDX-License-Identifier: Apache-2.0
"""BEVFormer-small in torch fp32: the deployed ``forward_trt`` graph (DerryHub/BEVFormer_tensorrt, the export behind
Autoware's ``bevformer_small.onnx``), built from :class:`~.weights.BEVFormerWeights` with the backbone BatchNorms
folded exactly where the port folds them. Ported from the verified research reference
(``research/bevformer/scripts/bevformer_ref.py``; research/bevformer/SPEC.md sections 4 and 6).

Two entry points with the same outputs ``(bev_embed [22500, 1, 256], outputs_classes [6, 1, 900, 10],
outputs_coords [6, 1, 900, 10])``:

- :meth:`BEVFormerNet.forward` takes the five ONNX inputs (``image``, ``prev_bev``, ``use_prev_bev``, ``can_bus``,
  ``lidar2img``) and computes the geometry INSIDE the graph in fp32, as the deployed graph does: the oracle form,
  compared with ONNX Runtime on the derived pure ONNX (``tests/test_reference_cpu.py``).
- :meth:`BEVFormerNet.forward_port` takes the geometry the TT port computes on the HOST (``host.geometry``:
  ``ref_cam`` / SCA tables from ``lidar2img`` in float64, the TSA shift, the prev_bev rotation index map) and runs the
  port's exact rewrites: the rotation as a row gather with a zero row, the spatial cross-attention REBATCHED per
  camera into ``SCA_CAPACITY`` rows and scatter-added back by a gather-form table (no dense 6 x 22,500 pass). Proven
  equal to :meth:`forward` on CPU before any device work (``tests/test_reference_host.py``).

Per-module taps (``self.taps = {}`` before a call; ``None`` disables them): ``backbone.stem``, ``backbone.layer1-4``,
``neck``, ``enc.can_bus_mlp``, ``enc.bev_queries``, ``enc.value``, ``enc.prev_bev_rot``, ``enc.shift``, ``ref_cam``,
``bev_mask_n``, ``encoder.layer{i}.{tsa,norm0,sca,norm1,ffn}``, ``encoder.layer{i}`` (layer output),
``decoder.layer{i}.{self_attn,norm0,cross_attn,norm1,ffn,ref}``, ``decoder.layer{i}``, ``head.cls{i}``,
``head.reg{i}``. The research goldens use the same names for the ones they hold.
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F

from . import config as C
from . import ops
from .weights import BEVFormerWeights

__all__ = ["BEVFormerNet", "PortGeometry", "reference_points_3d", "ref_2d", "point_sampling", "can_bus_shift",
           "inverse_sigmoid_mmdet", "inverse_sigmoid_dec"]

E = C.EMBED
Q = C.NUM_QUERY_BEV


# ------------------------------------------------------------------------------------------------ geometry (fp32)

def reference_points_3d() -> torch.Tensor:
    """``[1, 4, Q, 3]`` normalised (x, y, z) of the 4 pillar points of every BEV cell (DH encoder.py:169-195)."""
    Z = C.PC_RANGE[5] - C.PC_RANGE[2]
    zs = torch.linspace(0.5, Z - 0.5, C.PILLAR_POINTS).view(-1, 1, 1).expand(C.PILLAR_POINTS, C.BEV_H, C.BEV_W) / Z
    xs = torch.linspace(0.5, C.BEV_W - 0.5, C.BEV_W).view(1, 1, C.BEV_W).expand(C.PILLAR_POINTS, C.BEV_H, C.BEV_W)
    ys = torch.linspace(0.5, C.BEV_H - 0.5, C.BEV_H).view(1, C.BEV_H, 1).expand(C.PILLAR_POINTS, C.BEV_H, C.BEV_W)
    return torch.stack((xs / C.BEV_W, ys / C.BEV_H, zs), -1).reshape(1, C.PILLAR_POINTS, -1, 3)


def ref_2d() -> torch.Tensor:
    """``[1, Q, 1, 2]`` normalised BEV cell centres (x, y), row-major over (y, x)."""
    return reference_points_3d()[0, 0, :, :2].reshape(1, -1, 1, 2)


def point_sampling(ref_3d: torch.Tensor, lidar2img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """``point_sampling_trt`` (DH encoder.py:197-259) in fp32 -> ``ref_cam [6, Q, 4, 2]``, ``bev_mask_n [6, Q, 1]``
    (the any-of-4 camera mask normalised by the number of cameras that see the cell, clamp 1e-4)."""
    pr = C.PC_RANGE
    rp = ref_3d * torch.tensor([pr[3] - pr[0], pr[4] - pr[1], pr[5] - pr[2]]).view(1, 1, 1, 3) + torch.tensor(pr[:3])
    rp = torch.cat((rp, torch.ones_like(rp[..., :1])), -1).view(C.PILLAR_POINTS, 1, 1, -1, 4, 1)
    l2i = lidar2img.view(1, 1, C.NUM_CAMS, 1, 4, 4)
    cam = torch.matmul(l2i, rp).squeeze(-1)                                      # [4, 1, 6, Q, 4]
    eps = C.MSDA_EPS
    mask = (cam[..., 2:3] > eps).float()
    uv = cam[..., 0:2] / torch.maximum(cam[..., 2:3], torch.ones_like(cam[..., 2:3]) * eps)
    uv = uv / torch.tensor([C.IMG_HW[1], C.IMG_HW[0]], dtype=uv.dtype)
    mask = mask * (uv[..., 1:2] > 0.0) * (uv[..., 1:2] < 1.0) * (uv[..., 0:1] < 1.0) * (uv[..., 0:1] > 0.0)
    ref_cam = uv.permute(2, 1, 3, 0, 4).reshape(C.NUM_CAMS, -1, C.PILLAR_POINTS, 2)
    m = (1 - (1 - mask).prod(0)).view(C.NUM_CAMS, -1, 1)
    return ref_cam, m / torch.clamp(m.sum(0, keepdim=True), min=1e-4)


def can_bus_shift(can_bus: torch.Tensor) -> torch.Tensor:
    """TSA reference shift ``[1, 2]`` (x, y) in BEV units from can_bus[0], [1], [16], with DerryHub's atan2 surrogate
    (``get_bev_features_trt``, DH transformer.py:262-294), in fp32 as the deployed graph."""
    dx, dy = can_bus[0:1], can_bus[1:2]
    ego_angle = can_bus[16:17] / math.pi * 180
    tl = torch.sqrt(dx ** 2 + dy ** 2)
    ta = (torch.atan(dy / (dx + 1e-8)) + ((1 - torch.sign(dx)) / 2) * torch.sign(dy) * math.pi) / math.pi * 180
    ba = ego_angle - ta
    sy = tl * torch.cos(ba / 180 * math.pi) / C.GRID_LENGTH / C.BEV_H
    sx = tl * torch.sin(ba / 180 * math.pi) / C.GRID_LENGTH / C.BEV_W
    return torch.cat([sx, sy]).unsqueeze(0)


def inverse_sigmoid_mmdet(x: torch.Tensor, eps: float = C.INV_SIG_EPS) -> torch.Tensor:
    """mmdet ``inverse_sigmoid`` (the head's coordinate decode)."""
    x = x.clamp(min=0, max=1)
    return torch.log(x.clamp(min=eps) / (1 - x).clamp(min=eps))


def inverse_sigmoid_dec(x: torch.Tensor, eps: float = C.INV_SIG_EPS) -> torch.Tensor:
    """The decoder's reference update (DH decoder.py:24-37)."""
    x = x.clamp(min=eps, max=1 - eps)
    return torch.log(x / (1 - x))


@dataclass
class PortGeometry:
    """The per-frame host geometry of the TT port (``host.geometry`` builds it; all torch, fp32 / int64):

    ``ref_cam [6, Q, 4, 2]``; ``sca_rows [6, CAP]`` (BEV query index of every rebatched row, ``Q`` = padding);
    ``sca_counts [6]``; ``sca_gather [Q, K]`` (row of the stacked ``[6 * CAP + 1]`` SCA output table for each of the
    up to K cameras that see a query, ``6 * CAP`` = the zero row); ``sca_scale [Q, 1]`` (1 / cameras seeing it, 0 if
    none); ``shift [1, 2]``; ``rot_index [Q]`` (source row of the rotated prev_bev, ``Q`` = the zero row)."""

    ref_cam: torch.Tensor
    sca_rows: torch.Tensor
    sca_counts: torch.Tensor
    sca_gather: torch.Tensor
    sca_scale: torch.Tensor
    shift: torch.Tensor
    rot_index: torch.Tensor

    @classmethod
    def from_numpy(cls, d: Dict[str, np.ndarray]) -> "PortGeometry":
        def t(name, dtype):
            return torch.from_numpy(np.ascontiguousarray(d[name])).to(dtype)

        return cls(t("ref_cam", torch.float32), t("sca_rows", torch.int64), t("sca_counts", torch.int64),
                   t("sca_gather", torch.int64), t("sca_scale", torch.float32), t("shift", torch.float32),
                   t("rot_index", torch.int64))


# ------------------------------------------------------------------------------------------------------- model

class BEVFormerNet:
    """The network as torch fp32 functions over the checkpoint tensors (BN folded in the backbone)."""

    def __init__(self, weights: BEVFormerWeights, *, q_chunk: int = 4096):
        self.weights = weights
        self.q_chunk = int(q_chunk)
        sd = weights.sd
        self.p: Dict[str, torch.Tensor] = {k: torch.from_numpy(np.array(sd[k], np.float32)) for k in sd
                                           if not k.startswith("img_backbone") and sd[k].dtype != np.int64}
        self.conv: Dict[str, Tuple[torch.Tensor, torch.Tensor, int, int]] = {}
        for module, _, _, _, _ in C.backbone_conv_specs():
            cp = weights.conv(module)
            self.conv[module] = (torch.from_numpy(cp.weight), torch.from_numpy(cp.bias), cp.stride[0],
                                 cp.padding[0])
        h = "pts_bbox_head.positional_encoding"
        x_embed, y_embed = self.p[f"{h}.col_embed.weight"], self.p[f"{h}.row_embed.weight"]
        pos = torch.cat([x_embed.unsqueeze(0).expand(C.BEV_H, C.BEV_W, -1),
                         y_embed.unsqueeze(1).expand(C.BEV_H, C.BEV_W, -1)], dim=-1)
        self.bev_pos = pos.reshape(Q, E).contiguous()                              # [Q, 256], constant
        self.taps: Optional[Dict[str, torch.Tensor]] = None

    # ---- helpers ------------------------------------------------------------------------------------------------
    def _tap(self, name: str, value: torch.Tensor) -> None:
        if self.taps is not None:
            self.taps[name] = value.detach().clone()

    def _lin(self, name: str, x: torch.Tensor) -> torch.Tensor:
        return F.linear(x, self.p[f"{name}.weight"], self.p[f"{name}.bias"])

    def _ln(self, name: str, x: torch.Tensor) -> torch.Tensor:
        return F.layer_norm(x, (E,), self.p[f"{name}.weight"], self.p[f"{name}.bias"], eps=C.LN_EPS)

    def _conv(self, module: str, x: torch.Tensor, relu: bool) -> torch.Tensor:
        w, b, s, pad = self.conv[module]
        y = F.conv2d(x, w, b, stride=s, padding=pad)
        return F.relu(y) if relu else y

    # ---- image branch -------------------------------------------------------------------------------------------
    def _dcn(self, module: str, x: torch.Tensor) -> torch.Tensor:
        """DCNv2 conv2 + folded bn2 + ReLU: offset conv (3x3 -> 27: dy/dx of 9 taps as o1 | o2, mask logits),
        sigmoid mask, deformable sampling + GEMM with the BN-folded weight and bias."""
        out = self._conv(f"{module}.conv_offset", x, False)
        o1, o2, m = torch.chunk(out, 3, dim=1)
        w, b, _, _ = self.conv[module]
        return F.relu(ops.mdcn(x, torch.cat((o1, o2), dim=1), torch.sigmoid(m), w, b))

    def backbone(self, x: torch.Tensor) -> torch.Tensor:
        """ResNet-101 (caffe style, DCNv2 in stages 3-4) on ``[6, 3, 736, 1280]`` -> C5 ``[6, 2048, 23, 40]``."""
        y = F.max_pool2d(self._conv("img_backbone.conv1", x, True), 3, 2, 1)
        self._tap("backbone.stem", y)
        for s, (n, dcn) in enumerate(zip(C.RESNET_LAYERS, C.RESNET_DCN)):
            for b in range(n):
                p = f"img_backbone.layer{s + 1}.{b}"
                t = self._conv(f"{p}.conv1", y, True)
                t = self._dcn(f"{p}.conv2", t) if dcn else self._conv(f"{p}.conv2", t, True)
                t = self._conv(f"{p}.conv3", t, False)
                idt = self._conv(f"{p}.downsample.0", y, False) if b == 0 else y
                y = F.relu(t + idt)
            self._tap(f"backbone.layer{s + 1}", y)
        return y

    def neck(self, c5: torch.Tensor) -> torch.Tensor:
        """FPN on C5 only: lateral 1x1 (2048 -> 256) + 3x3 (256 -> 256), biases, no activation."""
        n = "img_neck"
        lat = F.conv2d(c5, self.p[f"{n}.lateral_convs.0.conv.weight"], self.p[f"{n}.lateral_convs.0.conv.bias"])
        out = F.conv2d(lat, self.p[f"{n}.fpn_convs.0.conv.weight"], self.p[f"{n}.fpn_convs.0.conv.bias"], padding=1)
        self._tap("neck", out)
        return out

    def extract_feat(self, image: torch.Tensor) -> torch.Tensor:
        return self.neck(self.backbone(image.reshape(C.NUM_CAMS, 3, C.IMG_HW[0], C.IMG_HW[1])))

    # ---- encoder pieces -----------------------------------------------------------------------------------------
    def bev_queries(self, can_bus: torch.Tensor) -> torch.Tensor:
        """``bev_embedding + can_bus_mlp(can_bus)`` -> ``[Q, 256]`` (Linear 18 -> 128, ReLU, Linear -> 256, ReLU,
        LayerNorm; one vector broadcast to every query)."""
        t = "pts_bbox_head.transformer.can_bus_mlp"
        cb = F.relu(self._lin(f"{t}.2", F.relu(self._lin(f"{t}.0", can_bus.view(1, 18)))))
        cb = self._ln(f"{t}.norm", cb)
        self._tap("enc.can_bus_mlp", cb.view(E))
        q = self.p["pts_bbox_head.bev_embedding.weight"] + cb
        self._tap("enc.bev_queries", q)
        return q

    def encoder_value(self, feat: torch.Tensor) -> torch.Tensor:
        """``[6, 256, 23, 40]`` -> SCA value ``[6, 920, 256]`` (+ cams_embeds[cam] + level_embeds[0])."""
        t = "pts_bbox_head.transformer"
        n, c, fh, fw = feat.shape
        f = feat.reshape(n, c, fh * fw).permute(0, 2, 1)
        v = f + self.p[f"{t}.cams_embeds"].view(n, 1, E) + self.p[f"{t}.level_embeds"][0].view(1, 1, E)
        self._tap("enc.value", v)
        return v

    def tsa(self, i: int, query: torch.Tensor, value2: torch.Tensor, hybrid_ref: torch.Tensor) -> torch.Tensor:
        """TemporalSelfAttention (+ residual): ``query [Q, 256]``, ``value2 [2, Q, 256]`` (queue: [prev; current]),
        ``hybrid_ref [2, Q, 1, 2]``."""
        a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.0"
        q = query + self.bev_pos
        qcat = torch.cat([value2[0], q], -1).view(1, Q, 2 * E)                        # [1, Q, 512]
        v = self._lin(f"{a}.value_proj", value2).view(C.TSA_QUEUE, Q, C.NUM_HEADS, C.HEAD_DIM)
        off = self._lin(f"{a}.sampling_offsets", qcat).view(1, Q, C.NUM_HEADS, C.TSA_QUEUE, 1, C.TSA_POINTS, 2)
        off = off.permute(0, 3, 1, 2, 4, 5, 6).reshape(C.TSA_QUEUE, Q, C.NUM_HEADS, C.TSA_POINTS * 2)
        aw = self._lin(f"{a}.attention_weights", qcat).view(1, Q, C.NUM_HEADS, C.TSA_QUEUE, 1, C.TSA_POINTS)
        aw = aw.permute(0, 3, 1, 2, 4, 5).reshape(C.TSA_QUEUE, Q, C.NUM_HEADS, C.TSA_POINTS)
        out = ops.msda(v, [(C.BEV_H, C.BEV_W)], hybrid_ref, off, aw, self.q_chunk).flatten(2)   # [2, Q, 256]
        return self._lin(f"{a}.output_proj", out.mean(0)) + query

    def _sca_proj(self, i: int, value: torch.Tensor, q_rows: torch.Tensor):
        a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.1.deformable_attention"
        v = self._lin(f"{a}.value_proj", value).view(C.NUM_CAMS, C.FEAT_LEN, C.NUM_HEADS, C.HEAD_DIM)
        off = self._lin(f"{a}.sampling_offsets", q_rows)
        aw = self._lin(f"{a}.attention_weights", q_rows)
        return v, off, aw

    def sca_dense(self, i: int, query: torch.Tensor, value: torch.Tensor, ref_cam: torch.Tensor,
                  mask_n: torch.Tensor) -> torch.Tensor:
        """SpatialCrossAttention as the TRT graph runs it: dense over all 6 x 22,500 (camera, query) pairs, then the
        normalised-mask sum over cameras (+ output_proj + residual)."""
        qr = query.unsqueeze(0).expand(C.NUM_CAMS, Q, E)
        v, off, aw = self._sca_proj(i, value, qr)
        ref = ref_cam.reshape(C.NUM_CAMS, Q, 1, C.PILLAR_POINTS * 2)
        out = ops.msda(v, [C.FEAT_HW], ref, off.view(C.NUM_CAMS, Q, C.NUM_HEADS, C.SCA_POINTS * 2),
                       aw.view(C.NUM_CAMS, Q, C.NUM_HEADS, C.SCA_POINTS), self.q_chunk).flatten(2)
        slots = (out * mask_n).sum(0)
        a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.1"
        return self._lin(f"{a}.output_proj", slots) + query

    def sca_rebatched(self, i: int, query: torch.Tensor, value: torch.Tensor, geo: PortGeometry) -> torch.Tensor:
        """The port's SCA: per camera, gather its ``SCA_CAPACITY`` rows (valid queries + padding), run the linears
        and the MSDA on those rows only, stack the six outputs with a zero row, and sum every query's rows through
        the gather table, times 1 / cameras (identical to :meth:`sca_dense`: a query seen by no camera gets 0)."""
        cap = geo.sca_rows.shape[1]
        q_pad = torch.cat([query, torch.zeros(1, E, dtype=query.dtype)], 0)              # row Q = padding
        rows = q_pad[geo.sca_rows.reshape(-1)].view(C.NUM_CAMS, cap, E)
        v, off, aw = self._sca_proj(i, value, rows)
        ref_pad = torch.cat([geo.ref_cam, torch.zeros(C.NUM_CAMS, 1, C.PILLAR_POINTS, 2)], 1)
        ref = torch.stack([ref_pad[c, geo.sca_rows[c]] for c in range(C.NUM_CAMS)])     # [6, cap, 4, 2]
        out = ops.msda(v, [C.FEAT_HW], ref.reshape(C.NUM_CAMS, cap, 1, C.PILLAR_POINTS * 2),
                       off.view(C.NUM_CAMS, cap, C.NUM_HEADS, C.SCA_POINTS * 2),
                       aw.view(C.NUM_CAMS, cap, C.NUM_HEADS, C.SCA_POINTS), self.q_chunk).flatten(2)
        table = torch.cat([out.reshape(C.NUM_CAMS * cap, E), torch.zeros(1, E, dtype=out.dtype)], 0)
        slots = table[geo.sca_gather].sum(1) * geo.sca_scale                             # [Q, 256]
        a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.1"
        return self._lin(f"{a}.output_proj", slots) + query

    def ffn(self, prefix: str, x: torch.Tensor) -> torch.Tensor:
        return x + self._lin(f"{prefix}.ffns.0.layers.1", F.relu(self._lin(f"{prefix}.ffns.0.layers.0.0", x)))

    def encoder(self, bev_q: torch.Tensor, value: torch.Tensor, prev_rot: torch.Tensor, use_prev_bev: torch.Tensor,
                shift: torch.Tensor, sca) -> torch.Tensor:
        """3 x (TSA -> LN -> SCA -> LN -> FFN -> LN). ``sca(i, query) -> query`` is the dense or rebatched SCA.
        The TSA value is ``use * [rot(prev_bev); bev_queries] + (1 - use) * [query; query]`` with the CURRENT
        layer's query in the second term (DH encoder.py:305, 587-589)."""
        use = use_prev_bev.reshape(())
        r2 = ref_2d()
        hybrid = torch.cat([r2 + shift.view(1, 1, 1, 2) * use, r2], 0)                     # [2, Q, 1, 2]
        prev = torch.stack([prev_rot, bev_q])                                               # [2, Q, 256]
        query = bev_q
        for i in range(C.ENC_LAYERS):
            pre = f"pts_bbox_head.transformer.encoder.layers.{i}"
            value2 = use * prev + (1 - use) * query.unsqueeze(0).expand(2, Q, E)
            query = self.tsa(i, query, value2, hybrid)
            self._tap(f"encoder.layer{i}.tsa", query)
            query = self._ln(f"{pre}.norms.0", query)
            self._tap(f"encoder.layer{i}.norm0", query)
            query = sca(i, query)
            self._tap(f"encoder.layer{i}.sca", query)
            query = self._ln(f"{pre}.norms.1", query)
            self._tap(f"encoder.layer{i}.norm1", query)
            query = self.ffn(pre, query)
            self._tap(f"encoder.layer{i}.ffn", query)
            query = self._ln(f"{pre}.norms.2", query)
            self._tap(f"encoder.layer{i}", query)
        return query

    # ---- decoder + heads ----------------------------------------------------------------------------------------
    def _mha(self, i: int, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
        """mmcv MultiheadAttention (q = k = x + pos, v = x, 8 heads, batch_first False) + residual."""
        a = f"pts_bbox_head.transformer.decoder.layers.{i}.attentions.0.attn"
        wq, wk, wv = self.p[f"{a}.in_proj_weight"].chunk(3, 0)
        bq, bk, bv = self.p[f"{a}.in_proj_bias"].chunk(3, 0)
        qk = x + pos
        n = x.shape[0]
        q = F.linear(qk, wq, bq).view(n, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1)
        k = F.linear(qk, wk, bk).view(n, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1)
        v = F.linear(x, wv, bv).view(n, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1)
        att = torch.softmax(torch.matmul(q, k.transpose(1, 2)) * (1.0 / math.sqrt(C.HEAD_DIM)), dim=-1)
        o = torch.matmul(att, v).transpose(0, 1).reshape(n, E)
        return x + self._lin(f"{a}.out_proj", o)

    def _cross(self, i: int, x: torch.Tensor, pos: torch.Tensor, bev: torch.Tensor, ref_xy: torch.Tensor):
        """CustomMSDeformableAttention over the BEV (150 x 150, 4 points at ``ref_xy``) + residual."""
        a = f"pts_bbox_head.transformer.decoder.layers.{i}.attentions.1"
        q = (x + pos).view(1, -1, E)
        n = q.shape[1]
        v = self._lin(f"{a}.value_proj", bev.view(1, Q, E)).view(1, Q, C.NUM_HEADS, C.HEAD_DIM)
        off = self._lin(f"{a}.sampling_offsets", q).view(1, n, C.NUM_HEADS, C.DEC_POINTS * 2)
        aw = self._lin(f"{a}.attention_weights", q).view(1, n, C.NUM_HEADS, C.DEC_POINTS)
        out = ops.msda(v, [(C.BEV_H, C.BEV_W)], ref_xy.view(1, n, 1, 2), off, aw, self.q_chunk).flatten(2)
        return x + self._lin(f"{a}.output_proj", out.view(n, E))

    def _cls(self, i: int, h: torch.Tensor) -> torch.Tensor:
        b = f"pts_bbox_head.cls_branches.{i}"
        h = F.relu(self._ln(f"{b}.1", self._lin(f"{b}.0", h)))
        h = F.relu(self._ln(f"{b}.4", self._lin(f"{b}.3", h)))
        return self._lin(f"{b}.6", h)

    def _reg(self, i: int, h: torch.Tensor) -> torch.Tensor:
        b = f"pts_bbox_head.reg_branches.{i}"
        h = F.relu(self._lin(f"{b}.2", F.relu(self._lin(f"{b}.0", h))))
        return self._lin(f"{b}.4", h)

    def decode(self, bev: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """6 decoder layers (MHA -> LN -> deformable cross-attention -> LN -> FFN -> LN, box refinement) + the six
        cls / reg heads -> ``outputs_classes``, ``outputs_coords`` ``[6, 1, 900, 10]`` (cx, cy, cz in metres)."""
        qe = self.p["pts_bbox_head.query_embedding.weight"]
        pos, x = qe[:, :E], qe[:, E:]
        ref = torch.sigmoid(self._lin("pts_bbox_head.transformer.reference_points", pos))   # [900, 3]
        init_ref = ref
        cls_all, crd_all = [], []
        pr = C.PC_RANGE
        for i in range(C.DEC_LAYERS):
            pre = f"pts_bbox_head.transformer.decoder.layers.{i}"
            ref_in = ref
            x = self._mha(i, x, pos)
            self._tap(f"decoder.layer{i}.self_attn", x)
            x = self._ln(f"{pre}.norms.0", x)
            self._tap(f"decoder.layer{i}.norm0", x)
            x = self._cross(i, x, pos, bev, ref[:, :2])
            self._tap(f"decoder.layer{i}.cross_attn", x)
            x = self._ln(f"{pre}.norms.1", x)
            self._tap(f"decoder.layer{i}.norm1", x)
            x = self.ffn(pre, x)
            self._tap(f"decoder.layer{i}.ffn", x)
            x = self._ln(f"{pre}.norms.2", x)
            self._tap(f"decoder.layer{i}", x)
            t = self._reg(i, x)                                                   # the head's reg == the update's
            ref = torch.sigmoid(torch.cat([t[:, 0:2] + inverse_sigmoid_dec(ref[:, :2]),
                                           t[:, 4:5] + inverse_sigmoid_dec(ref[:, 2:3])], -1))
            self._tap(f"decoder.layer{i}.ref", ref)
            r = inverse_sigmoid_mmdet(init_ref if i == 0 else ref_in)
            cls = self._cls(i, x)
            xy = torch.sigmoid(t[:, 0:2] + r[:, 0:2])
            z = torch.sigmoid(t[:, 4:5] + r[:, 2:3])
            crd = torch.cat([xy[:, 0:1] * (pr[3] - pr[0]) + pr[0], xy[:, 1:2] * (pr[4] - pr[1]) + pr[1], t[:, 2:4],
                             z * (pr[5] - pr[2]) + pr[2], t[:, 5:]], -1)
            self._tap(f"head.cls{i}", cls)
            self._tap(f"head.reg{i}", t)
            cls_all.append(cls)
            crd_all.append(crd)
        return torch.stack(cls_all).unsqueeze(1), torch.stack(crd_all).unsqueeze(1)

    # ---- entry points -------------------------------------------------------------------------------------------
    @torch.no_grad()
    def forward(self, image, prev_bev, use_prev_bev, can_bus, lidar2img):
        """The deployed graph: the five ONNX inputs (any array-likes, float32) -> the three ONNX outputs (torch)."""
        image, prev_bev, use_prev_bev, can_bus, lidar2img = (torch.as_tensor(np.asarray(a, np.float32)) for a in
                                                             (image, prev_bev, use_prev_bev, can_bus, lidar2img))
        feat = self.extract_feat(image)
        value = self.encoder_value(feat)
        bev_q = self.bev_queries(can_bus)
        shift = can_bus_shift(can_bus)
        self._tap("enc.shift", shift.view(2))
        pb = prev_bev.reshape(C.BEV_H, C.BEV_W, E).permute(2, 0, 1)
        pb = ops.rotate(pb, can_bus[17:18], C.ROTATE_CENTER).permute(1, 2, 0).reshape(Q, E)
        self._tap("enc.prev_bev_rot", pb)
        ref_cam, mask_n = point_sampling(reference_points_3d(), lidar2img)
        self._tap("ref_cam", ref_cam)
        self._tap("bev_mask_n", mask_n)
        bev = self.encoder(bev_q, value, pb, use_prev_bev, shift,
                           lambda i, q: self.sca_dense(i, q, value, ref_cam, mask_n))
        cls, crd = self.decode(bev)
        return bev.view(Q, 1, E), cls, crd

    @torch.no_grad()
    def forward_port(self, image, prev_bev, use_prev_bev, can_bus, geo: PortGeometry):
        """The TT port's decomposition with host geometry ``geo`` (see :class:`PortGeometry`)."""
        image, prev_bev, use_prev_bev, can_bus = (torch.as_tensor(np.asarray(a, np.float32)) for a in
                                                  (image, prev_bev, use_prev_bev, can_bus))
        feat = self.extract_feat(image)
        value = self.encoder_value(feat)
        bev_q = self.bev_queries(can_bus)
        self._tap("enc.shift", geo.shift.view(2))
        table = torch.cat([prev_bev.reshape(Q, E), torch.zeros(1, E)], 0)
        pb = table[geo.rot_index]                                                           # row gather, zero row
        self._tap("enc.prev_bev_rot", pb)
        self._tap("ref_cam", geo.ref_cam)
        bev = self.encoder(bev_q, value, pb, use_prev_bev, geo.shift,
                           lambda i, q: self.sca_rebatched(i, q, value, geo))
        cls, crd = self.decode(bev)
        return bev.view(Q, 1, E), cls, crd

    def run(self, *args: Any, taps: bool = False, port: bool = False):
        """``forward`` / ``forward_port`` returning numpy plus the taps (numpy) when ``taps``."""
        self.taps = {} if taps else None
        try:
            outs = (self.forward_port if port else self.forward)(*args)
            tp = {k: v.numpy() for k, v in (self.taps or {}).items()}
        finally:
            self.taps = None
        return tuple(o.numpy() for o in outs), tp