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

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

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

trellis_sparse_mlx.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

trellis_sparse_mlx.egg-info/requires.txt

@@ -0,0 +1,2 @@
+mlx>=0.20
+numpy>=1.24

trellis_sparse_mlx.egg-info/top_level.txt

@@ -0,0 +1 @@
+trellis_sparse_mlx

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",

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 (