From 15a5e744349ab50f89a7e3873b922921daa10976 Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 10:35:22 +1000 Subject: [PATCH] Split VarLenTensor / SparseTensor to match Pixal3D's hierarchy 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. --- trellis_sparse_mlx.egg-info/PKG-INFO | 8 +++ trellis_sparse_mlx.egg-info/SOURCES.txt | 13 ++++ .../dependency_links.txt | 1 + trellis_sparse_mlx.egg-info/requires.txt | 2 + trellis_sparse_mlx.egg-info/top_level.txt | 1 + trellis_sparse_mlx/__init__.py | 4 +- trellis_sparse_mlx/tensor.py | 62 +++++++++++++++---- 7 files changed, 78 insertions(+), 13 deletions(-) create mode 100644 trellis_sparse_mlx.egg-info/PKG-INFO create mode 100644 trellis_sparse_mlx.egg-info/SOURCES.txt create mode 100644 trellis_sparse_mlx.egg-info/dependency_links.txt create mode 100644 trellis_sparse_mlx.egg-info/requires.txt create mode 100644 trellis_sparse_mlx.egg-info/top_level.txt diff --git a/trellis_sparse_mlx.egg-info/PKG-INFO b/trellis_sparse_mlx.egg-info/PKG-INFO new file mode 100644 index 0000000..87e49a6 --- /dev/null +++ b/trellis_sparse_mlx.egg-info/PKG-INFO @@ -0,0 +1,8 @@ +Metadata-Version: 2.4 +Name: trellis-sparse-mlx +Version: 0.1.0 +Summary: MLX implementation of the TRELLIS-lineage sparse module (SubMConv3d et al) for Apple Silicon +License: MIT +Requires-Python: >=3.10 +Requires-Dist: mlx>=0.20 +Requires-Dist: numpy>=1.24 diff --git a/trellis_sparse_mlx.egg-info/SOURCES.txt b/trellis_sparse_mlx.egg-info/SOURCES.txt new file mode 100644 index 0000000..fb3d470 --- /dev/null +++ b/trellis_sparse_mlx.egg-info/SOURCES.txt @@ -0,0 +1,13 @@ +README.md +pyproject.toml +tests/test_ops.py +tests/test_sparse.py +trellis_sparse_mlx/__init__.py +trellis_sparse_mlx/conv.py +trellis_sparse_mlx/ops.py +trellis_sparse_mlx/tensor.py +trellis_sparse_mlx.egg-info/PKG-INFO +trellis_sparse_mlx.egg-info/SOURCES.txt +trellis_sparse_mlx.egg-info/dependency_links.txt +trellis_sparse_mlx.egg-info/requires.txt +trellis_sparse_mlx.egg-info/top_level.txt \ No newline at end of file diff --git a/trellis_sparse_mlx.egg-info/dependency_links.txt b/trellis_sparse_mlx.egg-info/dependency_links.txt new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/trellis_sparse_mlx.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/trellis_sparse_mlx.egg-info/requires.txt b/trellis_sparse_mlx.egg-info/requires.txt new file mode 100644 index 0000000..e2fcf5f --- /dev/null +++ b/trellis_sparse_mlx.egg-info/requires.txt @@ -0,0 +1,2 @@ +mlx>=0.20 +numpy>=1.24 diff --git a/trellis_sparse_mlx.egg-info/top_level.txt b/trellis_sparse_mlx.egg-info/top_level.txt new file mode 100644 index 0000000..2a93c19 --- /dev/null +++ b/trellis_sparse_mlx.egg-info/top_level.txt @@ -0,0 +1 @@ +trellis_sparse_mlx diff --git a/trellis_sparse_mlx/__init__.py b/trellis_sparse_mlx/__init__.py index 9abb36c..d257c37 100644 --- a/trellis_sparse_mlx/__init__.py +++ b/trellis_sparse_mlx/__init__.py @@ -14,7 +14,7 @@ than vendoring its own copy. """ from .conv import SubMConv3d, build_indice_map -from .tensor import SparseTensor, downsample, subdivide +from .tensor import SparseTensor, VarLenTensor, downsample, subdivide from .ops import ( LayerNorm32, SparseFeedForwardNet, @@ -29,7 +29,7 @@ from .ops import ( ) __all__ = [ - "SparseTensor", "subdivide", "downsample", + "SparseTensor", "VarLenTensor", "subdivide", "downsample", "SubMConv3d", "build_indice_map", "SparseLinear", "LayerNorm32", "SparseGroupNorm32", "SparseSiLU", "SparseGELU", "SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention", diff --git a/trellis_sparse_mlx/tensor.py b/trellis_sparse_mlx/tensor.py index 6261a01..f214bb2 100644 --- a/trellis_sparse_mlx/tensor.py +++ b/trellis_sparse_mlx/tensor.py @@ -18,16 +18,55 @@ import mlx.core as mx import numpy as np -class SparseTensor: - """N non-empty voxels, each with a coordinate and a feature vector.""" +class VarLenTensor: + """Variable-length batched features: `feats` plus an explicit slice-per-batch layout. - __slots__ = ("feats", "coords", "_scale", "_spatial_cache", "_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[int, int, int] = (1, 1, 1), + scale: Tuple = (1, 1, 1), spatial_cache: Optional[dict] = None, layout: Optional[List[slice]] = None, ): @@ -58,7 +97,11 @@ class SparseTensor: @property def layout(self) -> List[slice]: - """One slice per batch item. Relies on batch-contiguity, as upstream does.""" + """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) @@ -78,12 +121,9 @@ class SparseTensor: 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 + # -- 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 (