"""Sparse layer tests against torch. Unlike the submanifold conv — where spconv is uninstallable and the oracle had to be hand-written — every layer here has a real torch counterpart, so these compare against upstream's actual semantics rather than a paraphrase of them. The upstream forward bodies are reproduced verbatim (see modules/sparse/{norm,linear,nonlinearity}.py), including the [N_b,C] -> [1,C,N_b] GroupNorm reshape, which is the one that would fail silently if guessed. """ import sys from pathlib import Path import mlx.core as mx import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from lato_mlx.sparse.ops import ( # noqa: E402 LayerNorm32, SparseGroupNorm32, SparseMultiHeadAttention, SparseTransformerBlock, ) from lato_mlx.sparse.tensor import SparseTensor # noqa: E402 def make_batched(n_per_batch, channels, seed=0): """Batch-contiguous coords, as upstream requires.""" rng = np.random.default_rng(seed) coords, feats = [], [] for b, nb in enumerate(n_per_batch): for i in range(nb): coords.append((b, i // 16, (i // 4) % 4, i % 4)) feats.append(rng.standard_normal((nb, channels)).astype(np.float32)) return np.array(coords, dtype=np.int32), np.concatenate(feats, 0) def test_group_norm_matches_torch(groups=8, channels=32): n_per_batch = [37, 51] coords, feats = make_batched(n_per_batch, channels, seed=1) rng = np.random.default_rng(2) w = rng.standard_normal(channels).astype(np.float32) b = rng.standard_normal(channels).astype(np.float32) gn = SparseGroupNorm32(groups, channels) gn.weight, gn.bias = mx.array(w), mx.array(b) got = np.asarray(gn(SparseTensor(mx.array(feats), mx.array(coords))).feats) # upstream: per batch item, [N_b,C] -> permute -> [1,C,N_b] -> nn.GroupNorm tg = torch.nn.GroupNorm(groups, channels, eps=1e-5, affine=True) tg.weight.data = torch.tensor(w) tg.bias.data = torch.tensor(b) want = np.zeros_like(feats) off = 0 for nb in n_per_batch: bf = torch.tensor(feats[off : off + nb]) bf = bf.permute(1, 0).reshape(1, channels, -1) bf = tg(bf) want[off : off + nb] = bf.reshape(channels, -1).permute(1, 0).detach().numpy() off += nb err = np.abs(got - want).max() assert err < 2e-4, f"group norm err {err:.3g}" return err def test_group_norm_is_not_per_voxel(groups=8, channels=32): """Guard the easy-to-miss distinction: GroupNorm here is NOT a per-voxel norm.""" coords, feats = make_batched([40], channels, seed=5) gn = SparseGroupNorm32(groups, channels) got = np.asarray(gn(SparseTensor(mx.array(feats), mx.array(coords))).feats) per_voxel = torch.nn.functional.group_norm( torch.tensor(feats).reshape(40, channels), groups ).numpy() assert np.abs(got - per_voxel).max() > 1e-3, "matched per-voxel norm — reshape lost" return 0.0 def test_layer_norm_matches_torch(channels=64): rng = np.random.default_rng(3) feats = rng.standard_normal((50, channels)).astype(np.float32) ln = LayerNorm32(channels, affine=False, eps=1e-6) got = np.asarray(ln(mx.array(feats))) want = F.layer_norm(torch.tensor(feats), (channels,), eps=1e-6).numpy() err = np.abs(got - want).max() assert err < 1e-5, f"layer norm err {err:.3g}" return err def test_self_attention_matches_torch(channels=64, heads=8): n_per_batch = [23, 31] coords, feats = make_batched(n_per_batch, channels, seed=4) rng = np.random.default_rng(6) wq = rng.standard_normal((channels * 3, channels)).astype(np.float32) * 0.05 bq = rng.standard_normal((channels * 3,)).astype(np.float32) * 0.05 wo = rng.standard_normal((channels, channels)).astype(np.float32) * 0.05 bo = rng.standard_normal((channels,)).astype(np.float32) * 0.05 attn = SparseMultiHeadAttention(channels, heads) attn.to_qkv.weight, attn.to_qkv.bias = mx.array(wq), mx.array(bq) attn.to_out.weight, attn.to_out.bias = mx.array(wo), mx.array(bo) got = np.asarray(attn(SparseTensor(mx.array(feats), mx.array(coords))).feats) # reference: attention strictly within each batch item d = channels // heads want = np.zeros_like(feats) off = 0 for nb in n_per_batch: f = torch.tensor(feats[off : off + nb]) qkv = F.linear(f, torch.tensor(wq), torch.tensor(bq)).reshape(nb, 3, heads, d) q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] o = F.scaled_dot_product_attention( q.permute(1, 0, 2)[None], k.permute(1, 0, 2)[None], v.permute(1, 0, 2)[None] ) o = o[0].permute(1, 0, 2).reshape(nb, channels) want[off : off + nb] = ( F.linear(o, torch.tensor(wo), torch.tensor(bo)).detach().numpy() ) off += nb err = np.abs(got - want).max() assert err < 2e-4, f"attention err {err:.3g}" return err def test_attention_does_not_cross_batches(channels=32, heads=4): """Perturbing batch 1 must never change batch 0's output.""" coords, feats = make_batched([12, 12], channels, seed=7) attn = SparseMultiHeadAttention(channels, heads) a = np.asarray(attn(SparseTensor(mx.array(feats), mx.array(coords))).feats) f2 = feats.copy() f2[12:] += 10.0 b = np.asarray(attn(SparseTensor(mx.array(f2), mx.array(coords))).feats) assert np.abs(a[:12] - b[:12]).max() < 1e-5, "batch 0 changed — attention leaked" return 0.0 def test_transformer_block_shape(channels=64, heads=8): coords, feats = make_batched([20, 20], channels, seed=8) blk = SparseTransformerBlock(channels, heads) out = blk(SparseTensor(mx.array(feats), mx.array(coords))) assert out.feats.shape == (40, channels) assert np.isfinite(np.asarray(out.feats)).all(), "non-finite output" return 0.0 if __name__ == "__main__": tests = [ ("group norm vs torch", test_group_norm_matches_torch), ("group norm != per-voxel", test_group_norm_is_not_per_voxel), ("layer norm vs torch", test_layer_norm_matches_torch), ("self-attn vs torch", test_self_attention_matches_torch), ("attn batch isolation", test_attention_does_not_cross_batches), ("transformer block", test_transformer_block_shape), ] failed = 0 for name, fn in tests: try: err = fn() print(f" PASS {name:26s} (max err {err:.2e})") except AssertionError as e: print(f" FAIL {name:26s} {e}") failed += 1 except Exception as e: # noqa: BLE001 print(f" ERROR {name:26s} {type(e).__name__}: {e}") failed += 1 print(f"\n{len(tests)-failed}/{len(tests)} passed") sys.exit(1 if failed else 0)