monster/lato.2_mrp_mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MLX sparse core: SubMConv3d + SparseTensor + weight converter

Browse Source 
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.
This commit is contained in:
John 2026-08-02 10:04:24 +10:00
commit 97dcdfb54a
9 changed files with 693 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
.gitignore vendored Normal file
View File

@ -0,0 +1,7 @@
.venv/
ckpt/
weights/
upstream/
__pycache__/
*.pyc
bench/*.json

26
CLAUDE.md Normal file
View File

@ -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.

78
README.md Normal file
View File

@ -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.

0
lato_mlx/__init__.py Normal file
View File

159
lato_mlx/convert.py Normal file
View File

@ -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()

0
lato_mlx/sparse/__init__.py Normal file
View File

146
lato_mlx/sparse/conv.py Normal file
View File

@ -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)

115
lato_mlx/sparse/tensor.py Normal file
View File

@ -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),
)

162
tests/test_sparse.py Normal file
View File

@ -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)