cee423501c
Closes the op gap for Pixal3D. downsample now takes mode='max'|'mean'. This differs BETWEEN MODELS and is silent when wrong: LATO.2 hardcodes amax, Pixal3D parameterises it and defaults to mean. Callers must be explicit. upsample is a MASKED expansion, distinct from subdivide: subdivision.feats is an [N, factor**3] 0/1 mask, so a voxel emits 0-8 children rather than always 8. Pixal3D's decoder uses it to grow the occupied set selectively. Child-index decoding is subidx // factor**i % factor - axis 0 fastest-varying, the REVERSE of subdivide's C-order, which is exactly the kind of mixup that misplaces every child silently. Tests cover both modes differing, masked expansion landing at the right offsets, and a voxel with no flagged children contributing nothing.
208 lines
8.5 KiB
Python
208 lines
8.5 KiB
Python
"""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 trellis_sparse_mlx.conv import SubMConv3d, build_indice_map, _kernel_offsets
|
|
from trellis_sparse_mlx.tensor import SparseTensor, subdivide, downsample, upsample
|
|
|
|
|
|
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
|
|
|
|
|
|
def test_downsample_modes():
|
|
"""max vs mean is silent if confused: LATO.2 uses max, Pixal3D defaults to mean."""
|
|
c = np.array([[0, 0, 0, 0], [0, 0, 0, 1], [0, 2, 2, 2]], dtype=np.int32)
|
|
f = np.array([[1.0, 5.0], [3.0, 2.0], [9.0, 9.0]], dtype=np.float32)
|
|
x = SparseTensor(mx.array(f), mx.array(c))
|
|
mx_out = np.asarray(downsample(x, 2, "max").feats)
|
|
mean_out = np.asarray(downsample(x, 2, "mean").feats)
|
|
assert np.allclose(mx_out, [[3.0, 5.0], [9.0, 9.0]]), mx_out
|
|
assert np.allclose(mean_out, [[2.0, 3.5], [9.0, 9.0]]), mean_out
|
|
assert np.abs(mx_out - mean_out).max() > 0.1, "modes produced the same thing"
|
|
return 0.0
|
|
|
|
|
|
def test_masked_upsample():
|
|
"""Only flagged children are emitted, and each lands at the right offset."""
|
|
c = np.array([[0, 1, 1, 1]], dtype=np.int32)
|
|
f = np.array([[7.0]], dtype=np.float32)
|
|
# flag children 0 (offset 0,0,0) and 5 -> 5//1%2=1, 5//2%2=0, 5//4%2=1 -> (1,0,1)
|
|
sub = np.zeros((1, 8), dtype=np.float32)
|
|
sub[0, 0] = 1
|
|
sub[0, 5] = 1
|
|
out = upsample(SparseTensor(mx.array(f), mx.array(c)), SparseTensor(mx.array(sub), mx.array(c)), 2)
|
|
co = np.asarray(out.coords)
|
|
assert co.shape == (2, 4), co.shape
|
|
got = {tuple(r[1:]) for r in co.tolist()}
|
|
assert got == {(2, 2, 2), (3, 2, 3)}, got
|
|
assert np.allclose(np.asarray(out.feats), [[7.0], [7.0]]), "feats not replicated"
|
|
return 0.0
|
|
|
|
|
|
def test_upsample_can_drop_a_voxel():
|
|
"""A voxel with no flagged children contributes nothing."""
|
|
c = np.array([[0, 0, 0, 0], [0, 1, 1, 1]], dtype=np.int32)
|
|
f = np.array([[1.0], [2.0]], dtype=np.float32)
|
|
sub = np.zeros((2, 8), dtype=np.float32)
|
|
sub[1, 3] = 1 # only the second voxel survives
|
|
out = upsample(SparseTensor(mx.array(f), mx.array(c)), SparseTensor(mx.array(sub), mx.array(c)), 2)
|
|
assert out.coords.shape[0] == 1, out.coords.shape
|
|
assert np.allclose(np.asarray(out.feats), [[2.0]]), np.asarray(out.feats)
|
|
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),
|
|
("downsample max vs mean", test_downsample_modes),
|
|
("masked upsample", test_masked_upsample),
|
|
("upsample drops voxels", test_upsample_can_drop_a_voxel),
|
|
]
|
|
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)
|