# 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)