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 (