trellis_sparse_mrp_mlx/trellis_sparse_mlx/tensor.py

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


def downsample(x: "SparseTensor", factor: int = 2) -> "SparseTensor":
    """Downsample by `factor`, reducing colliding voxels with MAX.

    Upstream's docstring says "average pooling" but the implementation passes
    reduce="amax" (the `reduce='mean'` line is commented out). Following the code,
    not the docstring — mean vs max here is numerically silent in shape and would
    quietly change every downsampled feature.

    Output coordinates come out sorted by the same packed code upstream sorts on, so
    rows stay batch-contiguous as SparseTensor requires.
    """
    import numpy as _np

    c = _np.asarray(x.coords, dtype=_np.int64).copy()
    c[:, 1:] //= factor

    maxs = c[:, 1:].max(axis=0) + 1
    # OFFSET = reversed cumprod, matching upstream's packing
    off = _np.array(
        [maxs[0] * maxs[1] * maxs[2], maxs[1] * maxs[2], maxs[2], 1], dtype=_np.int64
    )
    code = (c * off).sum(axis=1)

    uniq, inv = _np.unique(code, return_inverse=True)
    feats = _np.asarray(x.feats)
    out = _np.full((uniq.shape[0], feats.shape[1]), -_np.inf, dtype=feats.dtype)
    _np.maximum.at(out, inv, feats)

    new_coords = _np.stack(
        [
            uniq // off[0],
            (uniq // off[1]) % maxs[0],
            (uniq // off[2]) % maxs[1],
            uniq % maxs[2],
        ],
        axis=-1,
    ).astype(_np.int32)
    return SparseTensor(
        mx.array(out),
        mx.array(new_coords),
        scale=tuple(s * factor for s in x._scale),
    )