Sparse tensor

SparseTensor is the central value type. It binds batched lattice coordinates to feature rows and carries coordinate identity metadata so sparse relations can be cached and reused.

Coordinate rows use (batch, x, y, z) order. Feature rows use the same row order as the coordinate buffer. active_rows separates the static buffer capacity from the dynamic number of valid rows:

\[\texttt{coords} \in \mathbb{Z}^{N\times4},\quad \texttt{feats} \in \mathbb{R}^{N\times C},\quad 0 \le \texttt{active\_rows}[0] \le N.\]

Feature-only operations preserve coordinate identity. Row-changing operations construct a new coordinate key and active-row scalar.

Related pages

• Concept model: Sparse tensor model

• Coordinate ownership: Coordinate management

• Feature-preserving operations: Feature operations

• Row-changing sparse algebra: Sparse tensor algebra

• Backend route inputs: Backend path selection

class mlx_lattice.core.tensor.SparseTensor(coords, feats, stride=1, *, coord_key=None, coord_manager=None, batch_counts=None, active_rows=None)

Bases: object

Sparse feature tensor indexed by batched integer coordinates.

SparseTensor is the public sparse value container used by convolution, pooling, sparse algebra, point/voxel conversion, and mlx_lattice.nn modules. Coordinates have shape (N, 4) ordered as (batch, x, y, z). Features have shape (N, C) and share row order with coordinates.

The object also carries coordinate identity metadata. A CoordinateManager owns the coordinate array, a CoordinateMapKey names it, and active_rows records the number of valid rows inside the static buffer capacity. Feature-only operations preserve that identity; row-changing operations create a new coordinate key.

Parameters:

• coords (array) – Integer coordinate rows with shape (N, 4). CPU paths accept int32 or int64. Metal sparse kernels require int32.

• feats (array) – Feature matrix with shape (N, C).

• stride (int | Sequence[int]) – Spatial lattice stride for the coordinates. An integer expands to (s, s, s).

• coord_key (CoordinateMapKey | None) – Existing coordinate key to reuse. When supplied, coords must be the manager-owned coordinate array for that key.

• coord_manager (CoordinateManager | None) – Manager that owns coord_key or that receives newly inserted coordinates.

• batch_counts (Sequence[int] | None) – Optional number of rows per batch. Required by global pooling and batch-partitioned utilities.

• active_rows (array | None) – Optional int32 scalar array with shape (1,). This lets native builders use a fixed-capacity coordinate buffer while considering only the active prefix.

coords: array

feats: array

stride: tuple[int, int, int]

coord_key: CoordinateMapKey

coord_manager: CoordinateManager

batch_counts: tuple[int, ...] | None

active_rows: array

property capacity: int

Static row capacity of the sparse buffers.

property active_count: array

Lazy MLX scalar array containing the number of active sparse rows.

property channels: int

Number of feature channels per sparse row.

property shape: tuple[int, int]

Sparse buffer shape as (capacity, channels).

property dtype: Dtype

Feature dtype.

property batch_indices: array

Batch column from coords.

property batch_rows: tuple[array, ...]

Row indices for each batch using batch_counts metadata.

property decomposed_coordinates: tuple[array, ...]

Spatial coordinates split by batch.

property decomposed_features: tuple[array, ...]

Feature rows split by batch.

astype(dtype)

Return a tensor with features converted to dtype.

Coordinate identity, stride, batch metadata, and active-row metadata are preserved because only the feature matrix changes.

Return type:

SparseTensor

Parameters:

dtype (Dtype)

replace(*, coords=None, feats=None, stride=None)

Return a sparse tensor with selected fields replaced.

Replacing only feats preserves coordinate identity. Replacing coords or changing stride inserts the new coordinate buffer into the existing manager and drops stale batch metadata because the row support may have changed.

Return type:

SparseTensor

Parameters:

• coords (array | None)

• feats (array | None)

• stride (int | Sequence[int] | None)

reuse_coords_from(other)

Attach this tensor’s features to other’s coordinate identity.

The two tensors must already describe the same coordinate identity. This helper is useful when a feature computation produced a fresh SparseTensor wrapper but the caller wants the metadata object from another tensor.

Return type:

SparseTensor

Parameters:

other (SparseTensor)

same_coords(other)

Return whether two tensors share coordinate identity.

This is an identity check over manager/key ownership, not a row-wise equality check. Two separate coordinate arrays with equal values do not share identity until they are registered under the same manager/key.

Return type:

bool

Parameters:

other (SparseTensor)