97dcdfb54a
The blocker for LATO.2 on Apple Silicon is one op, not the whole setup.sh --all CUDA stack. Measured: 5 of 7 checkpoints are fully dense, and every SparseConv3d in the model is constructed stride=1/padding=None, which upstream dispatches to spconv's SubMConv3d. No strided or inverse sparse conv is ever instantiated. - SubMConv3d in pure MLX via a sorted-key indice map (27 lookups/voxel vectorised, cached per coordinate set the way spconv uses indice_key) - SparseTensor container + subdivide upsampling - Converter handles the 5-D layout collision: spconv KRSC [O,kz,ky,kx,I] vs torch Conv3d [O,I,kz,ky,kx]. Rank alone is ambiguous; misreading it silently mangles the voxel encoder. - 7/7 tests pass vs an independent naive reference, max err 3e-7. spconv has no Metal build so there is no upstream oracle; the reference shares no indexing code. Kernel orientation (feats[c+d] vs feats[c-d]) remains unverified and is silent when wrong; --flip-kernel builds the mirror for an end-to-end A/B.
116 lines
4.3 KiB
Python
116 lines
4.3 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 SparseTensor:
|
||
"""N non-empty voxels, each with a coordinate and a feature vector."""
|
||
|
||
__slots__ = ("feats", "coords", "_scale", "_spatial_cache", "_layout")
|
||
|
||
def __init__(
|
||
self,
|
||
feats: mx.array,
|
||
coords: mx.array,
|
||
scale: Tuple[int, int, int] = (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. Relies on batch-contiguity, as upstream does."""
|
||
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 ----------------------------------------------------------------
|
||
def cache_get(self, key: str):
|
||
return self._spatial_cache.get(key)
|
||
|
||
def cache_put(self, key: str, value) -> None:
|
||
self._spatial_cache[key] = value
|
||
|
||
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),
|
||
)
|