70cd436eb9
Extracted from lato.2_mrp_mlx. The same sparse module underlies LATO.2, Pixal3D and the rest of the TRELLIS.2 family, and all of them are blocked on Apple Silicon by the same single op - submanifold conv - so it belongs in one tested package rather than vendored per port. 13/13 tests: 7 for the conv against a hand-written reference (spconv is uninstallable here so there is no upstream oracle), 6 for the remaining layers against torch. SubMConv3d runs 13.8ms at 128^3/128ch on m3ultra, 5.9x faster than the obvious per-offset loop.
172 lines
6.6 KiB
Python
172 lines
6.6 KiB
Python
"""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 trellis_sparse_mlx.ops import ( # noqa: E402
|
|
LayerNorm32,
|
|
SparseGroupNorm32,
|
|
SparseMultiHeadAttention,
|
|
SparseTransformerBlock,
|
|
)
|
|
from trellis_sparse_mlx.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)
|