lato.2_mrp_mlx/tests/test_ops.py

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