From 97dcdfb54a973a519f944c1fcd79cbcb8690cd3d Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 10:04:24 +1000 Subject: [PATCH] MLX sparse core: SubMConv3d + SparseTensor + weight converter 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. --- .gitignore | 7 ++ CLAUDE.md | 26 ++++++ README.md | 78 +++++++++++++++++ lato_mlx/__init__.py | 0 lato_mlx/convert.py | 159 +++++++++++++++++++++++++++++++++++ lato_mlx/sparse/__init__.py | 0 lato_mlx/sparse/conv.py | 146 ++++++++++++++++++++++++++++++++ lato_mlx/sparse/tensor.py | 115 +++++++++++++++++++++++++ tests/test_sparse.py | 162 ++++++++++++++++++++++++++++++++++++ 9 files changed, 693 insertions(+) create mode 100644 .gitignore create mode 100644 CLAUDE.md create mode 100644 README.md create mode 100644 lato_mlx/__init__.py create mode 100644 lato_mlx/convert.py create mode 100644 lato_mlx/sparse/__init__.py create mode 100644 lato_mlx/sparse/conv.py create mode 100644 lato_mlx/sparse/tensor.py create mode 100644 tests/test_sparse.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..34e0395 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +.venv/ +ckpt/ +weights/ +upstream/ +__pycache__/ +*.pyc +bench/*.json diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..ab023ce --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,26 @@ +# lato.2_mrp_mlx — working notes + +MLX port of LATO.2 (factorised mesh gen: vertex flow -> connectivity flow) for Apple Silicon. + +## Key facts established by measurement (don't re-derive) +- Upstream is CUDA-locked via `modules/sparse/__init__.py`: BACKEND accepts only + spconv/torchsparse, ATTN only xformers/flash_attn. No SDPA fallback exists upstream. +- The ONLY CUDA-locked op that matters is submanifold 3x3x3 conv. Every SparseConv3d in + the model uses default `stride=1, padding=None` -> SubMConv3d. No strided/inverse conv. +- Of 7 checkpoints, only vvae.pt has sparse conv kernels (18). vflow.pt is sparse-typed + but conv-free. voxel_encoder.pt is dense nn.Conv3d (29). Rest are fully dense. +- spconv weight layout is KRSC `[out, kz, ky, kx, in]` (verified against vvae.pt). + torch nn.Conv3d is `[out, in, kz, ky, kx]`. Both are 5-D -> `classify_5d()` disambiguates; + do not treat rank-5 as automatically sparse. +- Upsampling is `SparseSubdivide` (coords*2 + unit cube, feats replicated), not inverse conv. + +## Unresolved +- Kernel orientation (correlation vs convolution) is unverified and numerically silent. + Settle it via V-VAE reconstruction quality, not by argument. `--flip-kernel` builds the + mirrored weights. + +## Conventions +- Project venv at `.venv` (mlx + numpy + torch + trimesh). torch is CONVERSION-ONLY. +- `upstream/LATO.2` is vendored read-only reference; never edit it. +- Tests must compare against an independent implementation, not a second copy of the same + indexing logic. spconv is unavailable here so there is no upstream numerical oracle. diff --git a/README.md b/README.md new file mode 100644 index 0000000..5c54b2e --- /dev/null +++ b/README.md @@ -0,0 +1,78 @@ +# lato.2_mrp_mlx + +An MLX port of [LATO.2](https://github.com/LoHhhha/LATO.2) — factorised 3D mesh generation +(vertex flow, then connectivity flow) — so it runs natively on Apple Silicon. + +## Why + +LATO.2 generates meshes to a **controllable vertex budget**, which is the interesting part: +it sidesteps the generate-dense-then-decimate loop that TRELLIS-style pipelines force on +you. But upstream inherits TRELLIS.2's `setup.sh` and hard-requires CUDA: + +```python +# upstream modules/sparse/__init__.py +BACKEND = 'spconv' # accepts only ['spconv', 'torchsparse'] — both CUDA-only +ATTN = 'flash_attn' # accepts only ['xformers', 'flash_attn'] — both CUDA-only +``` + +No SDPA fallback, no MPS path. Neither sparse backend has a Metal build. + +## The scope, once measured + +The CUDA surface is far smaller than `setup.sh --all` implies. Of the seven released +checkpoints, five are entirely dense. Sparse code touches only `vertex_autoencoder` and +`vertex_structured_flow`, and **every** `SparseConv3d` in the model is constructed with the +defaults `stride=1, padding=None` — which upstream dispatches to `SubMConv3d`. + +So the whole blocker is one operation: **submanifold 3×3×3 convolution**. + +| Upstream | Here | +|---|---| +| `spconv.SubMConv3d` | `lato_mlx/sparse/conv.py` — gather/scatter over a sorted-key indice map | +| `SparseInverseConv3d` | never instantiated upstream; not needed | +| strided sparse conv | never used; upsampling is `SparseSubdivide` (coord expansion ×8) | +| `nn.Conv3d` (voxel encoder) | dense, maps to `mlx.nn.Conv3d` | +| `flash_attn` / `xformers` | `attn_mode="full"` everywhere → plain SDPA | + +## Status + +- [x] `SparseTensor` container + `subdivide` +- [x] `SubMConv3d` (k=3 and k=1) — **7/7 correctness tests pass**, max err 3e-7 vs an + independent naive reference +- [x] Weight converter, all 7 checkpoints → MLX safetensors (3.3 GB) +- [ ] Remaining sparse ops: `SparseLinear`, `SparseGroupNorm32`, activations, attention +- [ ] Model graphs: V-VAE, V-Flow, T-VAE, T-Flow, encoders +- [ ] End-to-end inference +- [ ] Fleet benchmark (m1max / m2max / m4pro / m1ultra / m3ultra) + +## Correctness + +`spconv` cannot be installed here — that is the reason this port exists — so there is no +numerical diff against upstream. Instead `tests/test_sparse.py` checks the vectorised +implementation against a deliberately naive one written straight from the definition +(dict lookup, per-voxel loop). They share no indexing code, so an off-by-one cannot hide +in both. Tests also cover batch isolation, isolated voxels, and indice-map hit rates. + +**One assumption remains unverified**: whether spconv gathers `feats[c+d]` +(cross-correlation — the deep-learning convention, and what this implements) or +`feats[c-d]`. A flipped kernel is numerically silent. It gets settled end-to-end: the +V-VAE is an autoencoder, so a clean reconstruction confirms the orientation. Run the +converter with `--flip-kernel` to test the alternative without touching code. + +## Use + +```bash +uv venv --python 3.12 .venv +VIRTUAL_ENV=.venv uv pip install mlx numpy torch trimesh + +hf download 0x4c48/LATO.2 --local-dir ckpt # 3.3 GB upstream weights +.venv/bin/python -m lato_mlx.convert --ckpt ckpt --out weights +.venv/bin/python tests/test_sparse.py +``` + +Upstream source is vendored read-only under `upstream/LATO.2` for reference. + +## Licence + +Upstream LATO.2 is MIT (Copyright the LATO.2 authors); its sparse module carries +Microsoft and VAST-AI-Research copyright, also MIT. This port is MIT on the same terms. diff --git a/lato_mlx/__init__.py b/lato_mlx/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/lato_mlx/convert.py b/lato_mlx/convert.py new file mode 100644 index 0000000..307c514 --- /dev/null +++ b/lato_mlx/convert.py @@ -0,0 +1,159 @@ +"""Convert LATO.2 PyTorch checkpoints to MLX safetensors. + +Run with any torch-bearing interpreter; torch is needed only to unpickle the .pt files +and is not a runtime dependency of the port itself. + + python -m lato_mlx.convert --ckpt ckpt --out weights + +Weight layouts +-------------- +spconv 2.x stores SubMConv3d weights in KRSC order, verified empirically against the +released vvae.pt (e.g. `encoder.downsample.0.conv1.conv.weight` is (64,3,3,3,32) = +[out, kz, ky, kx, in]). The MLX SubMConv3d here wants [K^3, in, out], so: + + [O, kz, ky, kx, I] --transpose--> [kz, ky, kx, I, O] --reshape--> [K^3, I, O] + +The C-order flatten of (kz,ky,kx) matches `_kernel_offsets`, which meshgrids +`arange(k) - k//2` with indexing="ij" — so kernel index 13 is the centre tap at k=3. + +UNVERIFIED ASSUMPTION — kernel orientation +------------------------------------------ +Whether spconv gathers `feats[c + d]` (cross-correlation, the deep-learning convention +and what this port implements) or `feats[c - d]` (true convolution) cannot be checked +numerically on this machine, because spconv has no Metal build — which is the whole +reason this port exists. A flipped kernel is silent: shapes and norms look right and the +output is subtly wrong. + +It is settled end-to-end rather than by assertion: the V-VAE is an autoencoder, so if +encode→decode reconstructs the input mesh, the orientation is right; a flip yields +visibly broken geometry. `--flip-kernel` writes the mirrored variant so the alternative +is a one-flag experiment rather than a code change. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np + +# Checkpoints built on the sparse module tree. Measured against the released weights: +# only vvae actually carries SubMConv3d kernels (18 of them); vflow is sparse in the +# structural sense (SparseTensor + SparseLinear + sparse attention) but has no conv. +SPARSE_CKPTS = {"vvae.pt", "vflow.pt"} + + +def _load_state_dict(path: Path): + import torch + + obj = torch.load(str(path), map_location="cpu", weights_only=False) + if not isinstance(obj, dict): + obj = obj.state_dict() + for key in ("state_dict", "model", "module"): + inner = obj.get(key) + if isinstance(inner, dict): + obj = inner + break + return obj + + +def classify_5d(name: str, shape: tuple) -> str: + """'krsc' (spconv sparse) vs 'conv3d' (dense torch) — they collide at 5-D. + + spconv SubMConv3d : [out, kz, ky, kx, in] -> dims 1..3 are the cubic kernel + torch nn.Conv3d : [out, in, kz, ky, kx] -> dims 2..4 are the cubic kernel + + LATO.2 contains both (the V-VAE uses sparse convs; the voxel encoder uses dense + nn.Conv3d), so guessing by rank alone silently mangles one of them. + """ + krsc = shape[1] == shape[2] == shape[3] + conv3d = shape[2] == shape[3] == shape[4] + if krsc and not conv3d: + return "krsc" + if conv3d and not krsc: + return "conv3d" + if krsc and conv3d: + # e.g. (O,3,3,3,3): genuinely ambiguous by shape. spconv weights always sit + # under the wrapper's `.conv.weight`; fail loudly rather than coin-flip. + if name.endswith(".conv.weight"): + return "krsc" + raise ValueError( + f"{name}: ambiguous 5-D layout {shape} — cannot tell KRSC from Conv3d" + ) + raise ValueError(f"{name}: unrecognised 5-D layout {shape}") + + +def convert_tensor(name: str, arr: np.ndarray, flip_kernel: bool) -> tuple: + """Returns (array, kind). Only spconv KRSC kernels are reshaped.""" + if arr.ndim != 5: + return arr, "dense" + kind = classify_5d(name, arr.shape) + if kind == "conv3d": + # Dense conv is handled by mlx.nn.Conv3d, which wants [out, kz, ky, kx, in] + # (channels-last), so move the input-channel axis to the end. + return arr.transpose(0, 2, 3, 4, 1).copy(), "conv3d" + o, kz, ky, kx, i = arr.shape + w = arr.transpose(1, 2, 3, 4, 0).reshape(kz * ky * kx, i, o) + if flip_kernel: + w = w[::-1].copy() # mirror all offsets: correlation <-> convolution + return w, "krsc" + + +def convert_checkpoint(src: Path, dst: Path, flip_kernel: bool = False) -> dict: + import mlx.core as mx + + sd = _load_state_dict(src) + out, n_krsc, n_conv3d, skipped = {}, 0, 0, 0 + for k, v in sd.items(): + if not hasattr(v, "detach"): + skipped += 1 + continue + arr = v.detach().to("cpu").float().numpy() + arr, kind = convert_tensor(k, arr, flip_kernel) + n_krsc += kind == "krsc" + n_conv3d += kind == "conv3d" + out[k] = arr + + dst.parent.mkdir(parents=True, exist_ok=True) + mx.save_safetensors(str(dst), {k: mx.array(v) for k, v in out.items()}) + return { + "tensors": len(out), + "sparse_kernels": n_krsc, + "dense_conv3d": n_conv3d, + "skipped": skipped, + "mb": dst.stat().st_size / 1e6, + } + + +def main() -> None: + ap = argparse.ArgumentParser(description="LATO.2 -> MLX weight converter") + ap.add_argument("--ckpt", default="ckpt", help="dir of upstream .pt files") + ap.add_argument("--out", default="weights", help="output dir for .safetensors") + ap.add_argument( + "--flip-kernel", + action="store_true", + help="mirror sparse kernel offsets (see module docstring: orientation experiment)", + ) + ap.add_argument("--only", default=None, help="convert just this file, e.g. vvae.pt") + a = ap.parse_args() + + src_dir, out_dir = Path(a.ckpt), Path(a.out) + files = sorted(src_dir.glob("*.pt")) + if a.only: + files = [f for f in files if f.name == a.only] + if not files: + raise SystemExit(f"no .pt files in {src_dir}") + + for f in files: + dst = out_dir / (f.stem + ".safetensors") + info = convert_checkpoint(f, dst, a.flip_kernel) + tag = " [sparse]" if f.name in SPARSE_CKPTS else "" + print( + f" {f.name:20s} -> {dst.name:26s} " + f"{info['tensors']:4d} tensors, {info['sparse_kernels']:2d} sparse, {info['dense_conv3d']:2d} conv3d, " + f"{info['mb']:7.1f} MB{tag}" + ) + + +if __name__ == "__main__": + main() diff --git a/lato_mlx/sparse/__init__.py b/lato_mlx/sparse/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/lato_mlx/sparse/conv.py b/lato_mlx/sparse/conv.py new file mode 100644 index 0000000..8f1d520 --- /dev/null +++ b/lato_mlx/sparse/conv.py @@ -0,0 +1,146 @@ +"""Submanifold sparse 3D convolution in pure MLX. + +This is the ONLY genuinely CUDA-locked operation in LATO.2's inference path. Upstream +routes it to `spconv.SubMConv3d` (or torchsparse); neither has a Metal build, which is +what has kept LATO.2 — and TRELLIS before it — off Apple Silicon. + +Every SparseConv3d in the LATO.2 model code is constructed with the defaults +`stride=1, padding=None`, which upstream dispatches to SubMConv3d. So only the +submanifold case is needed, at kernel sizes 3 and 1. + +Submanifold semantics: the output occupies EXACTLY the input coordinates (no dilation of +the occupied set). For each output voxel c: + + out[c] = bias + sum over kernel offsets d of W[d] @ feats[c + d] (c+d occupied) + +Absent neighbours contribute nothing. Implemented by gathering into a feature matrix with +one appended zero row, so "missing" is index N and needs no masking in the hot loop. + +Building the indice map is the expensive part and depends only on the coordinate set, so +it is cached on the SparseTensor under `indice_key` — the same trick spconv uses, and the +reason upstream threads `indice_key=f"res_{resolution}"` through the ResBlocks. +""" + +from __future__ import annotations + +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from .tensor import SparseTensor + +_MISSING = -1 + + +def _kernel_offsets(k: int) -> np.ndarray: + """Kernel offsets in C order over (dz, dy, dx), centred — matches spconv's ordering.""" + r = np.arange(k) - (k // 2) + return np.stack(np.meshgrid(r, r, r, indexing="ij"), axis=-1).reshape(-1, 3) + + +def build_indice_map(coords: mx.array, kernel_size: int) -> np.ndarray: + """[K^3, N] int32 — for each offset, the row of the neighbour, or -1 if unoccupied. + + Uses a sorted-key binary search rather than a Python dict: at k=3 this is 27 lookups + per voxel, and a per-voxel dict lookup would dominate runtime for any real mesh. + """ + c = np.asarray(coords, dtype=np.int64) + n = c.shape[0] + if n == 0: + return np.full((kernel_size**3, 0), _MISSING, dtype=np.int32) + + offsets = _kernel_offsets(kernel_size) + pad = kernel_size // 2 + + # Encode (batch,z,y,x) into one int64. Shift by `pad` so that neighbour coordinates + # of -1 stay non-negative and cannot alias onto a real cell at the opposite edge. + lo = c.min(axis=0) - pad + ext = (c.max(axis=0) + pad) - lo + 1 + strides = np.array( + [ext[1] * ext[2] * ext[3], ext[2] * ext[3], ext[3], 1], dtype=np.int64 + ) + + def encode(arr: np.ndarray) -> np.ndarray: + return ((arr - lo) * strides).sum(axis=1) + + keys = encode(c) + order = np.argsort(keys, kind="stable") + sorted_keys = keys[order] + + imap = np.empty((offsets.shape[0], n), dtype=np.int32) + for i, d in enumerate(offsets): + probe = c.copy() + probe[:, 1:] += d # batch index (column 0) never shifts + pk = encode(probe) + pos = np.searchsorted(sorted_keys, pk) + pos_clipped = np.clip(pos, 0, n - 1) + hit = sorted_keys[pos_clipped] == pk + imap[i] = np.where(hit, order[pos_clipped], _MISSING).astype(np.int32) + return imap + + +class SubMConv3d(nn.Module): + """Submanifold sparse conv. Weight layout [K^3, in_channels, out_channels].""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 3, + bias: bool = True, + indice_key: Optional[str] = None, + ): + super().__init__() + if kernel_size % 2 != 1: + raise ValueError(f"kernel_size must be odd, got {kernel_size}") + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.indice_key = indice_key + + scale = (in_channels * kernel_size**3) ** -0.5 + self.weight = mx.random.uniform( + -scale, scale, (kernel_size**3, in_channels, out_channels) + ) + if bias: + self.bias = mx.zeros((out_channels,)) + + def _indice_map(self, x: SparseTensor) -> np.ndarray: + key = f"imap_k{self.kernel_size}_{self.indice_key}" + cached = x.cache_get(key) + if cached is not None: + return cached + imap = build_indice_map(x.coords, self.kernel_size) + x.cache_put(key, imap) + return imap + + def __call__(self, x: SparseTensor) -> SparseTensor: + n = x.feats.shape[0] + + # k=1 touches only the centre voxel, so it is exactly a per-voxel linear — + # skip the indice map entirely. + if self.kernel_size == 1: + out = x.feats @ self.weight[0] + if hasattr(self, "bias"): + out = out + self.bias + return x.replace(out) + + imap = self._indice_map(x) + + # One appended zero row: absent neighbours index it and contribute nothing, + # which avoids a per-offset boolean mask in the accumulation loop. + feats_pad = mx.concatenate( + [x.feats, mx.zeros((1, self.in_channels), dtype=x.feats.dtype)], axis=0 + ) + idx = np.where(imap == _MISSING, n, imap) + + out = mx.zeros((n, self.out_channels), dtype=x.feats.dtype) + for i in range(imap.shape[0]): + gathered = mx.take(feats_pad, mx.array(idx[i]), axis=0) + out = out + gathered @ self.weight[i].astype(x.feats.dtype) + + if hasattr(self, "bias"): + out = out + self.bias.astype(x.feats.dtype) + return x.replace(out) diff --git a/lato_mlx/sparse/tensor.py b/lato_mlx/sparse/tensor.py new file mode 100644 index 0000000..fd5f376 --- /dev/null +++ b/lato_mlx/sparse/tensor.py @@ -0,0 +1,115 @@ +"""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), + ) diff --git a/tests/test_sparse.py b/tests/test_sparse.py new file mode 100644 index 0000000..942f5c7 --- /dev/null +++ b/tests/test_sparse.py @@ -0,0 +1,162 @@ +"""Correctness tests for the MLX sparse core. + +spconv cannot be installed on this machine — that is the entire reason this port exists — +so there is no way to diff against upstream numerically here. Instead the vectorised MLX +implementation is checked against a deliberately naive, obviously-correct reference +written straight from the definition of submanifold convolution (dict lookup, per-voxel +Python loop). The two share no indexing code, so an off-by-one in the fast path cannot +hide in both. +""" + +import sys +from pathlib import Path + +import mlx.core as mx +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from lato_mlx.sparse.conv import SubMConv3d, build_indice_map, _kernel_offsets +from lato_mlx.sparse.tensor import SparseTensor, subdivide + + +def reference_subm_conv(coords, feats, weight, bias, k): + """Definition of submanifold conv, written for obviousness, not speed.""" + occupied = {tuple(c): i for i, c in enumerate(coords.tolist())} + offsets = _kernel_offsets(k) + n, out_c = coords.shape[0], weight.shape[2] + out = np.zeros((n, out_c), dtype=np.float64) + for i, c in enumerate(coords.tolist()): + for oi, d in enumerate(offsets): + nb = (c[0], c[1] + d[0], c[2] + d[1], c[3] + d[2]) + j = occupied.get(nb) + if j is not None: + out[i] += feats[j].astype(np.float64) @ weight[oi].astype(np.float64) + if bias is not None: + out += bias.astype(np.float64) + return out + + +def random_sparse(n_vox, channels, batch=2, res=8, seed=0): + rng = np.random.default_rng(seed) + seen, coords = set(), [] + while len(coords) < n_vox: + b = int(rng.integers(0, batch)) + z, y, x = (int(v) for v in rng.integers(0, res, 3)) + if (b, z, y, x) in seen: + continue + seen.add((b, z, y, x)) + coords.append((b, z, y, x)) + coords.sort() # upstream requires batch-contiguous rows + coords = np.array(coords, dtype=np.int32) + feats = rng.standard_normal((n_vox, channels)).astype(np.float32) + return coords, feats + + +def test_subm_conv_matches_reference(k=3, n=180, cin=12, cout=7): + coords, feats = random_sparse(n, cin, seed=1) + rng = np.random.default_rng(2) + w = rng.standard_normal((k**3, cin, cout)).astype(np.float32) * 0.1 + b = rng.standard_normal((cout,)).astype(np.float32) + + conv = SubMConv3d(cin, cout, k, bias=True, indice_key="t") + conv.weight = mx.array(w) + conv.bias = mx.array(b) + + got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) + want = reference_subm_conv(coords, feats, w, b, k) + err = np.abs(got - want).max() + assert err < 2e-4, f"k={k} max abs err {err:.3g}" + return err + + +def test_kernel1_is_pointwise(n=64, cin=8, cout=5): + coords, feats = random_sparse(n, cin, seed=3) + rng = np.random.default_rng(4) + w = rng.standard_normal((1, cin, cout)).astype(np.float32) + conv = SubMConv3d(cin, cout, 1, bias=False, indice_key="p") + conv.weight = mx.array(w) + got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) + err = np.abs(got - feats @ w[0]).max() + assert err < 1e-4, f"k=1 err {err:.3g}" + return err + + +def test_isolated_voxel_sees_only_itself(): + """A voxel with no occupied neighbours must reduce to the centre tap alone.""" + coords = np.array([[0, 0, 0, 0], [0, 50, 50, 50]], dtype=np.int32) + feats = np.ones((2, 3), dtype=np.float32) + w = np.zeros((27, 3, 3), dtype=np.float32) + centre = 13 # index of (0,0,0) in centred C-order offsets + assert tuple(_kernel_offsets(3)[centre]) == (0, 0, 0) + w[centre] = np.eye(3) + conv = SubMConv3d(3, 3, 3, bias=False, indice_key="iso") + conv.weight = mx.array(w) + got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) + assert np.abs(got - feats).max() < 1e-6, got + return 0.0 + + +def test_batches_do_not_leak(): + """Same spatial cell in two batch items must not become neighbours.""" + coords = np.array([[0, 1, 1, 1], [1, 1, 1, 1]], dtype=np.int32) + feats = np.array([[1.0], [100.0]], dtype=np.float32) + w = np.ones((27, 1, 1), dtype=np.float32) # sum every occupied neighbour + conv = SubMConv3d(1, 1, 3, bias=False, indice_key="b") + conv.weight = mx.array(w) + got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) + assert np.allclose(got, [[1.0], [100.0]]), f"batch leak: {got}" + return 0.0 + + +def test_indice_map_hit_rate(): + """A fully dense block: interior voxels must find all 27 neighbours.""" + coords = np.array( + [[0, z, y, x] for z in range(4) for y in range(4) for x in range(4)], + dtype=np.int32, + ) + imap = build_indice_map(mx.array(coords), 3) + lin = {tuple(c): i for i, c in enumerate(coords.tolist())} + interior = [lin[(0, z, y, x)] for z in (1, 2) for y in (1, 2) for x in (1, 2)] + assert (imap[:, interior] != -1).all(), "interior voxel missing a neighbour" + corner = lin[(0, 0, 0, 0)] + assert (imap[:, corner] != -1).sum() == 8, "corner should see exactly 8 of 27" + return 0.0 + + +def test_subdivide(): + coords = np.array([[0, 1, 2, 3]], dtype=np.int32) + feats = np.array([[5.0, 6.0]], dtype=np.float32) + out = subdivide(SparseTensor(mx.array(feats), mx.array(coords))) + c = np.asarray(out.coords) + assert c.shape == (8, 4) and out.feats.shape == (8, 2) + assert set(map(tuple, c[:, 1:].tolist())) == { + (2 + a, 4 + b, 6 + d) for a in (0, 1) for b in (0, 1) for d in (0, 1) + } + assert np.abs(np.asarray(out.feats) - feats).max() < 1e-6 + return 0.0 + + +if __name__ == "__main__": + tests = [ + ("subm k=3 vs reference", lambda: test_subm_conv_matches_reference(3)), + ("subm k=5 vs reference", lambda: test_subm_conv_matches_reference(5, n=140)), + ("k=1 is pointwise", test_kernel1_is_pointwise), + ("isolated voxel", test_isolated_voxel_sees_only_itself), + ("batch isolation", test_batches_do_not_leak), + ("indice map hit rate", test_indice_map_hit_rate), + ("subdivide", test_subdivide), + ] + failed = 0 + for name, fn in tests: + try: + err = fn() + print(f" PASS {name:28s} (max err {err:.2e})") + except AssertionError as e: + print(f" FAIL {name:28s} {e}") + failed += 1 + except Exception as e: # noqa: BLE001 + print(f" ERROR {name:28s} {type(e).__name__}: {e}") + failed += 1 + print(f"\n{len(tests)-failed}/{len(tests)} passed") + sys.exit(1 if failed else 0)