5dd3f11d92
Completes the sparse op surface Pixal3D's decoders need. spatial2channel/channel2spatial are sparse space-to-depth and its inverse. The slot index is sum_i (coord[i] % f) * f**i - axis 0 FASTEST varying, the reverse of C order. decodes it the same way; using C order in one and this in the other would misplace every child while keeping all shapes valid. SparseResBlockC2S3d widens to out_channels*8 so channel2spatial can redistribute those channels across the 8 children, with a predicted subdiv mask deciding which children exist - so the occupied set grows selectively (32 voxels -> 122, not 256). Tests cover exact round-trip, zero-fill of unoccupied slots, and selective growth.
357 lines
13 KiB
Python
357 lines
13 KiB
Python
"""SparseTensor for MLX — a coords/feats pair, mirroring LATO.2's upstream container.
|
||
|
||
Upstream wraps either `spconv.SparseConvTensor` or `torchsparse.SparseTensor`; both are
|
||
CUDA-only, which is what blocks LATO.2 on Apple Silicon. Nothing about the *data* needs
|
||
CUDA — it is just coordinates plus features — so this is a plain MLX reimplementation.
|
||
|
||
Layout matches upstream exactly so weight conversion stays a straight mapping:
|
||
coords : int32 [N, 4] -> (batch, z, y, x)
|
||
feats : float [N, C]
|
||
and rows belonging to one batch item are contiguous (upstream asserts this).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import List, Optional, Tuple
|
||
|
||
import mlx.core as mx
|
||
import numpy as np
|
||
|
||
|
||
class VarLenTensor:
|
||
"""Variable-length batched features: `feats` plus an explicit slice-per-batch layout.
|
||
|
||
Pixal3D splits the abstraction this way — `VarLenTensor` is the sequence container and
|
||
`SparseTensor(VarLenTensor)` adds voxel coordinates. LATO.2 has only the combined
|
||
class. Modelling the split here lets one package serve both without either port
|
||
having to adapt at every call site.
|
||
"""
|
||
|
||
def __init__(self, feats: mx.array, layout: Optional[List[slice]] = None):
|
||
self.feats = feats
|
||
self._layout = layout
|
||
self._spatial_cache: dict = {}
|
||
|
||
@staticmethod
|
||
def layout_from_seqlen(seqlen: List[int]) -> List[slice]:
|
||
out, start = [], 0
|
||
for n in seqlen:
|
||
out.append(slice(start, start + n))
|
||
start += n
|
||
return out
|
||
|
||
@staticmethod
|
||
def from_tensor_list(tensors: List[mx.array]) -> "VarLenTensor":
|
||
lay = VarLenTensor.layout_from_seqlen([t.shape[0] for t in tensors])
|
||
return VarLenTensor(mx.concatenate(tensors, axis=0), lay)
|
||
|
||
@property
|
||
def layout(self) -> List[slice]:
|
||
if self._layout is None:
|
||
self._layout = [slice(0, self.feats.shape[0])]
|
||
return self._layout
|
||
|
||
# Pixal3D's spelling of the cache API; LATO.2 uses cache_get/cache_put. Both work.
|
||
def get_spatial_cache(self, key: str):
|
||
return self._spatial_cache.get(key)
|
||
|
||
def register_spatial_cache(self, key: str, value) -> None:
|
||
self._spatial_cache[key] = value
|
||
|
||
|
||
class SparseTensor(VarLenTensor):
|
||
"""N non-empty voxels, each with a coordinate and a feature vector."""
|
||
|
||
def __init__(
|
||
self,
|
||
feats: mx.array,
|
||
coords: mx.array,
|
||
scale: Tuple = (1, 1, 1),
|
||
spatial_cache: Optional[dict] = None,
|
||
layout: Optional[List[slice]] = None,
|
||
):
|
||
if feats.shape[0] != coords.shape[0]:
|
||
raise ValueError(
|
||
f"feats/coords length mismatch: {feats.shape[0]} vs {coords.shape[0]}"
|
||
)
|
||
if coords.ndim != 2 or coords.shape[1] != 4:
|
||
raise ValueError(f"coords must be [N, 4] (batch,z,y,x), got {coords.shape}")
|
||
self.feats = feats
|
||
self.coords = coords
|
||
self._scale = tuple(scale)
|
||
# Indice maps are expensive to build and identical for every conv that shares a
|
||
# coordinate set — upstream exploits this via spconv's `indice_key`. Same idea.
|
||
self._spatial_cache = spatial_cache if spatial_cache is not None else {}
|
||
self._layout = layout
|
||
|
||
# -- basics ---------------------------------------------------------------
|
||
@property
|
||
def shape(self) -> Tuple[int, int]:
|
||
return (self.batch_size, self.feats.shape[1])
|
||
|
||
@property
|
||
def batch_size(self) -> int:
|
||
if self.coords.shape[0] == 0:
|
||
return 0
|
||
return int(mx.max(self.coords[:, 0]).item()) + 1
|
||
|
||
@property
|
||
def layout(self) -> List[slice]:
|
||
"""One slice per batch item, derived from coords. Relies on batch-contiguity.
|
||
|
||
Overrides VarLenTensor.layout, which has no coordinates to derive from and so
|
||
treats the whole tensor as a single sequence.
|
||
"""
|
||
if self._layout is None:
|
||
b = np.asarray(self.coords[:, 0], dtype=np.int64)
|
||
counts = np.bincount(b, minlength=self.batch_size)
|
||
offs = np.cumsum(counts)
|
||
self._layout = [
|
||
slice(int(offs[i] - counts[i]), int(offs[i])) for i in range(len(counts))
|
||
]
|
||
return self._layout
|
||
|
||
def replace(self, feats: mx.array) -> "SparseTensor":
|
||
"""New tensor, same coordinates — so the indice-map cache stays valid."""
|
||
return SparseTensor(
|
||
feats,
|
||
self.coords,
|
||
scale=self._scale,
|
||
spatial_cache=self._spatial_cache,
|
||
layout=self._layout,
|
||
)
|
||
|
||
# -- cache (LATO.2 spelling; see VarLenTensor for Pixal3D's) ---------------
|
||
cache_get = VarLenTensor.get_spatial_cache
|
||
cache_put = VarLenTensor.register_spatial_cache
|
||
|
||
def __repr__(self) -> str:
|
||
return (
|
||
f"SparseTensor(N={self.coords.shape[0]}, C={self.feats.shape[1]}, "
|
||
f"batch={self.batch_size}, scale={self._scale})"
|
||
)
|
||
|
||
|
||
def subdivide(x: SparseTensor) -> SparseTensor:
|
||
"""Upsample ×2 by splitting each voxel into its 8 children (nearest-neighbour).
|
||
|
||
Mirrors upstream `SparseSubdivide`: coords are doubled then offset by the unit
|
||
cube, features are replicated. Child order is the C-order of nonzero(ones(2,2,2)),
|
||
matching upstream's `torch.nonzero`, so replicated features line up identically.
|
||
"""
|
||
n = x.coords.shape[0]
|
||
offsets = np.stack(np.meshgrid(*[np.arange(2)] * 3, indexing="ij"), -1).reshape(-1, 3)
|
||
offsets = np.concatenate([np.zeros((8, 1), dtype=offsets.dtype), offsets], axis=1)
|
||
|
||
base = np.asarray(x.coords, dtype=np.int32).copy()
|
||
base[:, 1:] *= 2
|
||
new_coords = (base[:, None, :] + offsets[None, :, :]).reshape(n * 8, 4)
|
||
|
||
new_feats = mx.repeat(x.feats, 8, axis=0)
|
||
return SparseTensor(
|
||
new_feats,
|
||
mx.array(new_coords, dtype=mx.int32),
|
||
scale=tuple(s * 2 for s in x._scale),
|
||
)
|
||
|
||
|
||
def downsample(
|
||
x: "SparseTensor", factor: int = 2, mode: str = "max"
|
||
) -> "SparseTensor":
|
||
"""Downsample by `factor`, reducing colliding voxels with `mode` ("max" or "mean").
|
||
|
||
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.
|
||
"""
|
||
import numpy as _np
|
||
|
||
c = _np.asarray(x.coords, dtype=_np.int64).copy()
|
||
c[:, 1:] //= factor
|
||
|
||
maxs = c[:, 1:].max(axis=0) + 1
|
||
# OFFSET = reversed cumprod, matching upstream's packing
|
||
off = _np.array(
|
||
[maxs[0] * maxs[1] * maxs[2], maxs[1] * maxs[2], maxs[2], 1], dtype=_np.int64
|
||
)
|
||
code = (c * off).sum(axis=1)
|
||
|
||
uniq, inv = _np.unique(code, return_inverse=True)
|
||
feats = _np.asarray(x.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(
|
||
[
|
||
uniq // off[0],
|
||
(uniq // off[1]) % maxs[0],
|
||
(uniq // off[2]) % maxs[1],
|
||
uniq % maxs[2],
|
||
],
|
||
axis=-1,
|
||
).astype(_np.int32)
|
||
return SparseTensor(
|
||
mx.array(out),
|
||
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),
|
||
)
|
||
|
||
|
||
def _subidx_of(coords_np, factor: int):
|
||
"""Slot index of each voxel within its factor^3 cell.
|
||
|
||
subidx = sum_i (coord[i] % factor) * factor**i — axis 0 is the FASTEST varying,
|
||
which is the reverse of C order. `upsample` and `channel2spatial` decode it back the
|
||
same way (`subidx // factor**i % factor`); using C order in one place and this in the
|
||
other misplaces every child while keeping all shapes valid.
|
||
"""
|
||
import numpy as _np
|
||
|
||
r = coords_np[:, 1:] % factor
|
||
return sum(r[:, i] * (factor**i) for i in range(3)).astype(_np.int64)
|
||
|
||
|
||
def spatial2channel(x: "SparseTensor", factor: int = 2) -> "SparseTensor":
|
||
"""Sparse space-to-depth: [N, C] at res R -> [N', C*factor^3] at res R/factor.
|
||
|
||
Voxels sharing a coarse cell are packed into disjoint channel slots; slots with no
|
||
voxel are ZERO. Registers the inverse mapping in the spatial cache so a later
|
||
`channel2spatial` can undo it without being handed a subdivision mask.
|
||
"""
|
||
import numpy as _np
|
||
|
||
c = _np.asarray(x.coords, dtype=_np.int64)
|
||
feats = _np.asarray(x.feats)
|
||
n, ch = feats.shape
|
||
f3 = factor**3
|
||
|
||
coarse = c.copy()
|
||
coarse[:, 1:] //= factor
|
||
maxs = coarse[:, 1:].max(axis=0) + 1
|
||
off = _np.array(
|
||
[maxs[0] * maxs[1] * maxs[2], maxs[1] * maxs[2], maxs[2], 1], dtype=_np.int64
|
||
)
|
||
code = (coarse * off).sum(axis=1)
|
||
uniq, idx = _np.unique(code, return_inverse=True)
|
||
subidx = _subidx_of(c, factor)
|
||
|
||
out = _np.zeros((uniq.shape[0] * f3, ch), dtype=feats.dtype)
|
||
out[idx * f3 + subidx] = feats
|
||
new_coords = _np.stack(
|
||
[uniq // off[0], (uniq // off[1]) % maxs[0], (uniq // off[2]) % maxs[1], uniq % maxs[2]],
|
||
axis=-1,
|
||
).astype(_np.int32)
|
||
|
||
res = SparseTensor(
|
||
mx.array(out.reshape(uniq.shape[0], ch * f3)),
|
||
mx.array(new_coords),
|
||
scale=tuple(s * factor for s in x._scale),
|
||
)
|
||
res.register_spatial_cache(f"channel2spatial_{factor}", (c.astype(_np.int32), idx, subidx))
|
||
return res
|
||
|
||
|
||
def channel2spatial(
|
||
x: "SparseTensor", subdivision: "SparseTensor" = None, factor: int = 2
|
||
) -> "SparseTensor":
|
||
"""Inverse of `spatial2channel`: [N, C*factor^3] -> [N', C] at finer resolution.
|
||
|
||
Uses the cached mapping when this follows a `spatial2channel`; otherwise a
|
||
`subdivision` mask ([N, factor^3], nonzero = keep) selects which children exist.
|
||
Unlike `upsample`, features are SPLIT across children rather than replicated: child
|
||
j takes channel slot j.
|
||
"""
|
||
import numpy as _np
|
||
|
||
cache = x.cache_get(f"channel2spatial_{factor}")
|
||
f3 = factor**3
|
||
feats = _np.asarray(x.feats)
|
||
ch = feats.shape[1] // f3
|
||
|
||
if cache is not None:
|
||
new_coords, idx, subidx = cache
|
||
new_coords = _np.asarray(new_coords, dtype=_np.int32)
|
||
else:
|
||
if subdivision is None:
|
||
raise ValueError(
|
||
"channel2spatial needs a subdivision mask, or to follow a spatial2channel"
|
||
)
|
||
sub = _np.asarray(subdivision.feats) > 0
|
||
n_leaf = sub.sum(axis=1)
|
||
idx = _np.repeat(_np.arange(sub.shape[0]), n_leaf)
|
||
subidx = _np.nonzero(sub)[1]
|
||
base = _np.asarray(x.coords, dtype=_np.int64).copy()
|
||
base[:, 1:] *= factor
|
||
new_coords = base[idx]
|
||
for i in range(3):
|
||
new_coords[:, i + 1] += subidx // (factor**i) % factor
|
||
new_coords = new_coords.astype(_np.int32)
|
||
|
||
flat = feats.reshape(feats.shape[0] * f3, ch)
|
||
new_feats = flat[idx * f3 + subidx]
|
||
return SparseTensor(
|
||
mx.array(new_feats),
|
||
mx.array(new_coords),
|
||
scale=tuple(s / factor for s in x._scale),
|
||
)
|