15a5e74434
Pixal3D refactored the container into a VarLenTensor base (feats + explicit slice-per-batch layout, no coords) with SparseTensor(VarLenTensor) adding coordinates. LATO.2 has only the combined class. Modelling the split here means one package serves both without either port adapting at every call site. Also exposes Pixal3D's cache spelling (get/register_spatial_cache) alongside LATO.2's (cache_get/cache_put). 13/13 tests unchanged.
200 lines
7.2 KiB
Python
200 lines
7.2 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) -> "SparseTensor":
|
||
"""Downsample by `factor`, reducing colliding voxels with MAX.
|
||
|
||
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.
|
||
|
||
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)
|
||
out = _np.full((uniq.shape[0], feats.shape[1]), -_np.inf, dtype=feats.dtype)
|
||
_np.maximum.at(out, inv, 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)
|
||
return SparseTensor(
|
||
mx.array(out),
|
||
mx.array(new_coords),
|
||
scale=tuple(s * factor for s in x._scale),
|
||
)
|