Add masked upsample and parameterise downsample reduction (16/16 tests)

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.

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:

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

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),
+    )