"""Correctness tests for the MLX sparse core. spconv cannot be installed on this machine — that is the entire reason this port exists — so there is no way to diff against upstream numerically here. Instead the vectorised MLX implementation is checked against a deliberately naive, obviously-correct reference written straight from the definition of submanifold convolution (dict lookup, per-voxel Python loop). The two share no indexing code, so an off-by-one in the fast path cannot hide in both. """ import sys from pathlib import Path import mlx.core as mx import numpy as np sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from lato_mlx.sparse.conv import SubMConv3d, build_indice_map, _kernel_offsets from lato_mlx.sparse.tensor import SparseTensor, subdivide def reference_subm_conv(coords, feats, weight, bias, k): """Definition of submanifold conv, written for obviousness, not speed.""" occupied = {tuple(c): i for i, c in enumerate(coords.tolist())} offsets = _kernel_offsets(k) n, out_c = coords.shape[0], weight.shape[2] out = np.zeros((n, out_c), dtype=np.float64) for i, c in enumerate(coords.tolist()): for oi, d in enumerate(offsets): nb = (c[0], c[1] + d[0], c[2] + d[1], c[3] + d[2]) j = occupied.get(nb) if j is not None: out[i] += feats[j].astype(np.float64) @ weight[oi].astype(np.float64) if bias is not None: out += bias.astype(np.float64) return out def random_sparse(n_vox, channels, batch=2, res=8, seed=0): rng = np.random.default_rng(seed) seen, coords = set(), [] while len(coords) < n_vox: b = int(rng.integers(0, batch)) z, y, x = (int(v) for v in rng.integers(0, res, 3)) if (b, z, y, x) in seen: continue seen.add((b, z, y, x)) coords.append((b, z, y, x)) coords.sort() # upstream requires batch-contiguous rows coords = np.array(coords, dtype=np.int32) feats = rng.standard_normal((n_vox, channels)).astype(np.float32) return coords, feats def test_subm_conv_matches_reference(k=3, n=180, cin=12, cout=7): coords, feats = random_sparse(n, cin, seed=1) rng = np.random.default_rng(2) w = rng.standard_normal((k**3, cin, cout)).astype(np.float32) * 0.1 b = rng.standard_normal((cout,)).astype(np.float32) conv = SubMConv3d(cin, cout, k, bias=True, indice_key="t") conv.weight = mx.array(w) conv.bias = mx.array(b) got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) want = reference_subm_conv(coords, feats, w, b, k) err = np.abs(got - want).max() assert err < 2e-4, f"k={k} max abs err {err:.3g}" return err def test_kernel1_is_pointwise(n=64, cin=8, cout=5): coords, feats = random_sparse(n, cin, seed=3) rng = np.random.default_rng(4) w = rng.standard_normal((1, cin, cout)).astype(np.float32) conv = SubMConv3d(cin, cout, 1, bias=False, indice_key="p") conv.weight = mx.array(w) got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) err = np.abs(got - feats @ w[0]).max() assert err < 1e-4, f"k=1 err {err:.3g}" return err def test_isolated_voxel_sees_only_itself(): """A voxel with no occupied neighbours must reduce to the centre tap alone.""" coords = np.array([[0, 0, 0, 0], [0, 50, 50, 50]], dtype=np.int32) feats = np.ones((2, 3), dtype=np.float32) w = np.zeros((27, 3, 3), dtype=np.float32) centre = 13 # index of (0,0,0) in centred C-order offsets assert tuple(_kernel_offsets(3)[centre]) == (0, 0, 0) w[centre] = np.eye(3) conv = SubMConv3d(3, 3, 3, bias=False, indice_key="iso") conv.weight = mx.array(w) got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) assert np.abs(got - feats).max() < 1e-6, got return 0.0 def test_batches_do_not_leak(): """Same spatial cell in two batch items must not become neighbours.""" coords = np.array([[0, 1, 1, 1], [1, 1, 1, 1]], dtype=np.int32) feats = np.array([[1.0], [100.0]], dtype=np.float32) w = np.ones((27, 1, 1), dtype=np.float32) # sum every occupied neighbour conv = SubMConv3d(1, 1, 3, bias=False, indice_key="b") conv.weight = mx.array(w) got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) assert np.allclose(got, [[1.0], [100.0]]), f"batch leak: {got}" return 0.0 def test_indice_map_hit_rate(): """A fully dense block: interior voxels must find all 27 neighbours.""" coords = np.array( [[0, z, y, x] for z in range(4) for y in range(4) for x in range(4)], dtype=np.int32, ) imap = build_indice_map(mx.array(coords), 3) lin = {tuple(c): i for i, c in enumerate(coords.tolist())} interior = [lin[(0, z, y, x)] for z in (1, 2) for y in (1, 2) for x in (1, 2)] assert (imap[:, interior] != -1).all(), "interior voxel missing a neighbour" corner = lin[(0, 0, 0, 0)] assert (imap[:, corner] != -1).sum() == 8, "corner should see exactly 8 of 27" return 0.0 def test_subdivide(): coords = np.array([[0, 1, 2, 3]], dtype=np.int32) feats = np.array([[5.0, 6.0]], dtype=np.float32) out = subdivide(SparseTensor(mx.array(feats), mx.array(coords))) c = np.asarray(out.coords) assert c.shape == (8, 4) and out.feats.shape == (8, 2) assert set(map(tuple, c[:, 1:].tolist())) == { (2 + a, 4 + b, 6 + d) for a in (0, 1) for b in (0, 1) for d in (0, 1) } assert np.abs(np.asarray(out.feats) - feats).max() < 1e-6 return 0.0 if __name__ == "__main__": tests = [ ("subm k=3 vs reference", lambda: test_subm_conv_matches_reference(3)), ("subm k=5 vs reference", lambda: test_subm_conv_matches_reference(5, n=140)), ("k=1 is pointwise", test_kernel1_is_pointwise), ("isolated voxel", test_isolated_voxel_sees_only_itself), ("batch isolation", test_batches_do_not_leak), ("indice map hit rate", test_indice_map_hit_rate), ("subdivide", test_subdivide), ] failed = 0 for name, fn in tests: try: err = fn() print(f" PASS {name:28s} (max err {err:.2e})") except AssertionError as e: print(f" FAIL {name:28s} {e}") failed += 1 except Exception as e: # noqa: BLE001 print(f" ERROR {name:28s} {type(e).__name__}: {e}") failed += 1 print(f"\n{len(tests)-failed}/{len(tests)} passed") sys.exit(1 if failed else 0)