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.
258 lines
9.4 KiB
Python
258 lines
9.4 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),
|
||
)
|