diff --git a/tests/test_sparse.py b/tests/test_sparse.py index 0c069e6..c6b645f 100644 --- a/tests/test_sparse.py +++ b/tests/test_sparse.py @@ -17,7 +17,7 @@ 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 +from trellis_sparse_mlx.tensor import SparseTensor, subdivide, downsample, upsample def reference_subm_conv(coords, feats, weight, bias, k): @@ -137,6 +137,48 @@ def test_subdivide(): 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)), @@ -146,6 +188,9 @@ if __name__ == "__main__": ("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: diff --git a/trellis_sparse_mlx/__init__.py b/trellis_sparse_mlx/__init__.py index d257c37..dedc117 100644 --- a/trellis_sparse_mlx/__init__.py +++ b/trellis_sparse_mlx/__init__.py @@ -14,7 +14,7 @@ than vendoring its own copy. """ from .conv import SubMConv3d, build_indice_map -from .tensor import SparseTensor, VarLenTensor, downsample, subdivide +from .tensor import SparseTensor, VarLenTensor, downsample, subdivide, upsample from .ops import ( LayerNorm32, SparseFeedForwardNet, @@ -29,7 +29,7 @@ from .ops import ( ) __all__ = [ - "SparseTensor", "VarLenTensor", "subdivide", "downsample", + "SparseTensor", "VarLenTensor", "subdivide", "downsample", "upsample", "SubMConv3d", "build_indice_map", "SparseLinear", "LayerNorm32", "SparseGroupNorm32", "SparseSiLU", "SparseGELU", "SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention", diff --git a/trellis_sparse_mlx/tensor.py b/trellis_sparse_mlx/tensor.py index f214bb2..4dd1504 100644 --- a/trellis_sparse_mlx/tensor.py +++ b/trellis_sparse_mlx/tensor.py @@ -155,13 +155,20 @@ def subdivide(x: SparseTensor) -> SparseTensor: ) -def downsample(x: "SparseTensor", factor: int = 2) -> "SparseTensor": - """Downsample by `factor`, reducing colliding voxels with MAX. +def downsample( + x: "SparseTensor", factor: int = 2, mode: str = "max" +) -> "SparseTensor": + """Downsample by `factor`, reducing colliding voxels with `mode` ("max" or "mean"). - Upstream's docstring says "average pooling" but the implementation passes - reduce="amax" (the `reduce='mean'` line is commented out). Following the code, - not the docstring — mean vs max here is numerically silent in shape and would - quietly change every downsampled feature. + The reduction differs BETWEEN MODELS and is numerically silent if you get it wrong — + identical shapes, quietly different features: + + LATO.2 hardcodes reduce="amax" (its docstring says "average pooling"; the + `reduce='mean'` line is commented out — follow the code, not the prose). + Pixal3D parameterises it as `mode: Literal['mean','max'] = 'mean'`, i.e. it + DEFAULTS TO MEAN. + + So callers must be explicit. LATO.2 passes mode="max". Output coordinates come out sorted by the same packed code upstream sorts on, so rows stay batch-contiguous as SparseTensor requires. @@ -180,8 +187,15 @@ def downsample(x: "SparseTensor", factor: int = 2) -> "SparseTensor": uniq, inv = _np.unique(code, return_inverse=True) feats = _np.asarray(x.feats) - out = _np.full((uniq.shape[0], feats.shape[1]), -_np.inf, dtype=feats.dtype) - _np.maximum.at(out, inv, feats) + if mode == "max": + out = _np.full((uniq.shape[0], feats.shape[1]), -_np.inf, dtype=feats.dtype) + _np.maximum.at(out, inv, feats) + elif mode == "mean": + out = _np.zeros((uniq.shape[0], feats.shape[1]), dtype=feats.dtype) + _np.add.at(out, inv, feats) + out /= _np.bincount(inv, minlength=uniq.shape[0])[:, None].astype(feats.dtype) + else: + raise ValueError(f"mode must be 'max' or 'mean', got {mode!r}") new_coords = _np.stack( [ @@ -197,3 +211,47 @@ def downsample(x: "SparseTensor", factor: int = 2) -> "SparseTensor": mx.array(new_coords), scale=tuple(s * factor for s in x._scale), ) + + +def upsample( + x: "SparseTensor", subdivision: "SparseTensor", factor: int = 2 +) -> "SparseTensor": + """Masked ×`factor` upsample — expand each voxel only into its FLAGGED children. + + Distinct from `subdivide`, which expands every voxel into all 8 children + unconditionally. Here `subdivision.feats` is a [N, factor**3] 0/1 mask saying which + children survive, so a voxel can produce anywhere from 0 to 8 outputs. Pixal3D's + decoder uses this to grow the occupied set selectively rather than blowing up ×8 at + every level. + + Child index decoding matches upstream: `subidx // factor**i % factor` for axis i, + i.e. axis 0 is the fastest-varying — note this is the REVERSE of `subdivide`'s + C-order, and mixing them up misplaces every child. + + Scale divides by `factor` (finer), where downsample multiplies. + """ + import numpy as _np + + sub = _np.asarray(subdivision.feats) + if sub.ndim != 2 or sub.shape[0] != x.coords.shape[0]: + raise ValueError( + f"subdivision must be [N, factor**3] matching N={x.coords.shape[0]}, " + f"got {sub.shape}" + ) + mask = sub > 0 + n_leaf = mask.sum(axis=1) + parent = _np.repeat(_np.arange(sub.shape[0]), n_leaf) + subidx = _np.nonzero(mask)[1] + + base = _np.asarray(x.coords, dtype=_np.int64).copy() + base[:, 1:] *= factor + new_coords = base[parent] + for i in range(3): + new_coords[:, i + 1] += subidx // (factor**i) % factor + + new_feats = mx.take(x.feats, mx.array(parent), axis=0) + return SparseTensor( + new_feats, + mx.array(new_coords.astype(_np.int32)), + scale=tuple(s / factor for s in x._scale), + )