code/tt_transfusion/tests/test_reference_rewrites_host.py

# SPDX-License-Identifier: Apache-2.0
"""The TT-friendly forms of ``reference/rewrites.py`` equal the literal modules of ``reference/model.py`` with the real
weights (CPU, torch fp32): what the device graph computes is proven here before device time is spent. Inputs are
seeded random tensors of the real channel counts (small spatial sizes), plus ties for the local max.

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_transfusion/tests/test_reference_rewrites_host.py
"""
from __future__ import annotations

import numpy as np
import pytest

pytest.importorskip("onnx")
torch = pytest.importorskip("torch")
F = torch.nn.functional

from tt_transfusion.reference import rewrites as rw  # noqa: E402
from tt_transfusion.reference.model import TransFusionTorch  # noqa: E402
from tt_transfusion.reference.weights import load_weights  # noqa: E402
from tt_transfusion.tests import _data  # noqa: E402

RNG = np.random.default_rng(20261008)


@pytest.fixture(scope="module")
def net():
    torch.set_num_threads(2)
    return TransFusionTorch(load_weights(_data.weights_dir()), threads=2)


def _t(a):
    return torch.from_numpy(np.array(a, dtype=np.float32, copy=True))


def _pillars(n=300, slots=20, feats=11):
    """Random decorated pillars: 1..20 points (a third of them full), padded slots zero."""
    num = RNG.integers(1, slots + 1, n)
    num[: n // 3] = slots
    f = RNG.normal(0.0, 1.0, (n, slots, feats)).astype(np.float32)
    f[:, :, 3] = RNG.integers(0, 256, (n, slots))                     # intensity scale
    f *= (np.arange(slots)[None, :] < num[:, None])[:, :, None]
    return f, num


def test_pfn_tile_form_and_32_slot_pillars(net):
    f, num = _pillars()
    p = net.weights.pfn
    with torch.no_grad():
        ref = net.pfn(f).numpy()
        f32 = rw.pad_pillar_slots(f)                                   # (n, 32, 16): slots 20-31 repeat slot 0
        w0 = np.zeros((rw.PFN_IN_PADDED, p.w0.shape[1]), np.float32)
        w0[: p.w0.shape[0]] = p.w0
        got = rw.pfn_tile_form(_t(f32), _t(w0), _t(p.b0), _t(p.w1a), _t(p.w1b), _t(p.b1)).numpy()
        # negative control: zero rows in slots 20-31 add relu(b0) to the maxima of the full pillars
        zero = np.zeros_like(f32)
        zero[:, :20, :11] = f
        bad = rw.pfn_tile_form(_t(zero), _t(w0), _t(p.b0), _t(p.w1a), _t(p.w1b), _t(p.b1)).numpy()
    assert f32.shape == (300, 32, 16) and np.array_equal(f32[:, 20:, :11], np.repeat(f[:, :1], 12, axis=1))
    np.testing.assert_allclose(got, ref, rtol=1e-5, atol=2e-5)
    full = num == 20
    assert np.abs(bad[full] - ref[full]).max() > 1e-3                 # the trap the duplicated slots avoid
    np.testing.assert_allclose(bad[~full], ref[~full], rtol=1e-5, atol=2e-5)


def test_k_split_shared_conv(net):
    x = [_t(np.maximum(RNG.normal(0, 1, (1, 128, 24, 24)), 0)) for _ in range(3)]
    shared = net.weights.head.shared
    with torch.no_grad():
        ref = net.layer("head.shared", torch.cat(x, dim=1))
        got = rw.k_split_torch(x, shared)
    np.testing.assert_allclose(got.numpy(), ref.numpy(), rtol=1e-4, atol=1e-4)


def test_deblock_rewrites(net):
    d0, _, d2 = net.weights.deblocks
    x0 = np.maximum(RNG.normal(0, 1, (1, 64, 16, 16)), 0).astype(np.float32)
    x2 = np.maximum(RNG.normal(0, 1, (1, 256, 8, 8)), 0).astype(np.float32)
    with torch.no_grad():
        ref0 = net.layer("neck.deblock0", _t(x0)).numpy()[0].transpose(1, 2, 0)        # (8, 8, 128)
        ref2 = net.layer("neck.deblock2", _t(x2)).numpy()[0].transpose(1, 2, 0)        # (16, 16, 128)
    wl, bl = rw.deblock0_as_linear(d0)
    s2d = rw.space_to_depth_nhwc(x0[0].transpose(1, 2, 0), 2)                           # (8, 8, 256)
    got0 = np.maximum(s2d.reshape(64, 256) @ wl + bl, 0).reshape(8, 8, 128)
    np.testing.assert_allclose(got0, ref0, rtol=1e-4, atol=1e-4)
    wl2, bl2 = rw.deconv_as_linear(d2)
    y = x2[0].transpose(1, 2, 0).reshape(64, 256) @ wl2 + bl2
    got2 = np.maximum(rw.depth_to_space_nhwc(y, 8, 8, 2, 128), 0)
    np.testing.assert_allclose(got2, ref2, rtol=1e-4, atol=1e-4)


def test_masked_local_max_chain_is_exact(net):
    """C22's chain with TransFusion's masks equals MaxPool p0 + zero border + equality bit for bit, with ties."""
    heat = np.round(RNG.uniform(0, 1, (1, 5, 192, 192)) * 64) / 64                    # many exact ties
    heat[0, 2, 50:60, 70:80] = 0.5                                                     # a flat plateau
    heat = _t(heat)
    lp = torch.zeros_like(heat)
    lp[:, :, 1:-1, 1:-1] = F.max_pool2d(heat, 3, 1, 0)
    literal = heat * (heat == lp).to(heat.dtype)
    m_pass, m_eq = rw.local_max_masks(192, 192, 5)
    got = rw.local_max_masked_torch(heat, m_pass, m_eq)
    assert torch.equal(got, literal)
    assert float(got[0, :, 0].abs().max()) == 0.0 and float(got[0, :, :, -1].abs().max()) == 0.0


def test_query_init_tables_are_exact(net):
    h = net.weights.head
    pos = RNG.integers(0, 36864, 500)
    assert np.array_equal(rw.query_pos_from_index(pos, 192), h.bev_pos[pos])
    onehot = np.zeros((1, 5, 500), np.float32)
    cls = RNG.integers(0, 5, 500)
    onehot[0, cls, np.arange(500)] = 1.0
    with torch.no_grad():
        conv = F.conv1d(_t(onehot), _t(h.class_encoding_w[:, :, None]), _t(h.class_encoding_b))[0].T.numpy()
    np.testing.assert_array_equal(h.class_table[cls], conv)


def test_cross_attention_constant_kv_tables_and_padded_heads(net):
    d = net.weights.head.decoder
    kpe = net.key_pos_embed().numpy()                                                  # (36864, 128)
    lidar = RNG.normal(0, 1, (36864, 128)).astype(np.float32)
    x = RNG.normal(0, 1, (500, 128)).astype(np.float32)
    with torch.no_grad():
        ref = net._mha("cross_attn", _t(x), _t(lidar + kpe)).numpy()
        ck, cv = rw.cross_attention_kv_tables(d.cross_attn, kpe)
        got = rw.padded_head_attention(_t(x), _t(lidar), _t(lidar), d.cross_attn, k_const=ck, v_const=cv).numpy()
        ref_self = net._mha("self_attn", _t(x), _t(x)).numpy()
        got_self = rw.padded_head_attention(_t(x), _t(x), _t(x), d.self_attn).numpy()
    np.testing.assert_allclose(got, ref, rtol=1e-4, atol=1e-4)
    np.testing.assert_allclose(got_self, ref_self, rtol=1e-4, atol=1e-4)


def test_merged_prediction_heads(net):
    x = RNG.normal(0, 1, (500, 128)).astype(np.float32)
    merged = rw.merged_prediction_head_weights(net.weights.head)
    wh, bh, wo, bo, slices = merged
    assert wh.shape == (128, 384) and wo.shape == (384, 32) and slices[-1] == ("heatmap", 10, 15)
    with torch.no_grad():
        got = rw.merged_heads_torch(_t(x), merged)
        y = _t(x).T[None]
        for name, hid, fin in net.weights.head.prediction_heads:
            ref = net.layer(fin.name, net.layer(hid.name, y))
            np.testing.assert_allclose(got[name].numpy(), ref.numpy(), rtol=1e-4, atol=1e-4)