Sparse tensor algebra¶
Sparse algebra functions combine tensors by coordinate identity or by value-aligned joins. If two tensors share coordinate identity, operations use a direct feature path. Otherwise the operation builds row maps:
where -1 means the coordinate is missing from that side of the join.
Join modes follow database-style sparse support semantics: inner keeps
coordinates present on both sides, left and right preserve one side,
and outer keeps the union.
SparseTensor also provides operator shortcuts for the common binary paths:
x + y maps to sparse_add, x - y maps to sparse_sub, and
x * y / x & y map to intersection-style sparse_mul. Use the named
functions when you need a non-default join or custom fill values.
Related pages¶
Algebra concept reference: Coordinate-aligned sparse algebra
Sparse tensor identity: Sparse tensor model
Coordinate alignment core: Coordinate utilities
Workflow guidance for branch joins: End-to-end sparse workflow
- mlx_lattice.ops.tensor.sparse_collate(coords, feats, *, stride=1)[source]¶
Collate unbatched sparse coordinates/features into one tensor.
Input coordinates must have shape
(N, 3); the function prepends a batch column and recordsbatch_countsmetadata. Feature arrays must each have shape(N_i, C)and all coordinate arrays must share dtype.
- mlx_lattice.ops.tensor.align_sparse(lhs, rhs, *, join='inner')[source]¶
Build coordinate-value alignment metadata for two sparse tensors.
Both tensors must have the same sparse stride and coordinate dtype. The returned row maps can be used to gather feature rows under
inner,left,right, oroutersparse support semantics.- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
join (SparseJoin)
- mlx_lattice.ops.tensor.gather_aligned_features(x, rows, *, fill=0.0)[source]¶
Gather sparse features with
-1rows filled byfill.rowsmust be anint32vector. Non-negative entries gatherx.feats; negative entries produce the scalar fill value in every channel.- Return type:
array- Parameters:
x (SparseTensor)
rows (array)
fill (float)
- mlx_lattice.ops.tensor.sparse_binary_op(lhs, rhs, op, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]¶
Apply an elementwise binary op after coordinate-value alignment.
If operands share coordinate identity, the operation runs directly on feature matrices. Otherwise it builds a sparse alignment and fills missing rows according to
lhs_fillandrhs_fill.- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
op (Literal['add', 'sub', 'mul', 'maximum', 'minimum'])
join (SparseJoin)
lhs_fill (float)
rhs_fill (float)
- mlx_lattice.ops.tensor.sparse_add(lhs, rhs, *, join='outer')[source]¶
Add sparse tensors after coordinate alignment.
- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
join (SparseJoin)
- mlx_lattice.ops.tensor.sparse_sub(lhs, rhs, *, join='outer')[source]¶
Subtract sparse tensors after coordinate alignment.
- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
join (SparseJoin)
- mlx_lattice.ops.tensor.sparse_mul(lhs, rhs, *, join='inner')[source]¶
Multiply sparse tensors after coordinate alignment.
- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
join (SparseJoin)
- mlx_lattice.ops.tensor.sparse_maximum(lhs, rhs, *, join='inner')[source]¶
Take elementwise maximum after coordinate alignment.
- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
join (SparseJoin)
- mlx_lattice.ops.tensor.sparse_minimum(lhs, rhs, *, join='inner')[source]¶
Take elementwise minimum after coordinate alignment.
- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
join (SparseJoin)
- mlx_lattice.ops.tensor.cat(tensors, *, join='inner')[source]¶
Concatenate sparse features, aligning coordinates when needed.
Multiple tensors are supported when they already share coordinate identity. Value-aligned concatenation currently accepts two tensors because the join support and fill behavior must be unambiguous.
- Return type:
- Parameters:
tensors (Sequence[SparseTensor])
join (SparseJoin)
- mlx_lattice.ops.tensor.sparse_cat_aligned(lhs, rhs, *, join='inner')[source]¶
Concatenate two sparse tensors after coordinate-value alignment.
- Return type:
- Parameters:
lhs (SparseTensor)
rhs (SparseTensor)
join (SparseJoin)
- mlx_lattice.ops.tensor.crop(x, *, min_coord, max_coord)[source]¶
Crop a sparse tensor to an inclusive spatial coordinate box.
min_coordandmax_coordare spatial triples overx, y, z. The batch column is not filtered.- Return type:
- Parameters:
x (SparseTensor)
- mlx_lattice.ops.tensor.replace_feature(x, feats)[source]¶
Return
xwith its feature matrix replaced.The replacement must satisfy the
SparseTensorrow contract: same row count asx.coordsand two-dimensional(N, C_new)shape.- Return type:
- Parameters:
x (SparseTensor)
feats (array)
- mlx_lattice.ops.tensor.prune(x, rows)[source]¶
Keep the specified sparse rows.
rowsis a one-dimensional index vector. The result contains gathered coordinates and features and receives a new coordinate key in the same manager.- Return type:
- Parameters:
x (SparseTensor)
rows (array)
- mlx_lattice.ops.tensor.prune_mask(x, mask)[source]¶
Keep sparse rows selected by a boolean mask.
maskmust have length equal to the sparse buffer capacity. Selected rows are gathered into a compact sparse tensor.- Return type:
- Parameters:
x (SparseTensor)
mask (array)
- mlx_lattice.ops.tensor.topk_rows(x, counts, *, rho=1.0)[source]¶
Select top-scoring rows per batch using the first feature channel.
countsgives the requested rows per batch andrhoscales those counts before clamping to the available rows in each batch. Requiresx.batch_counts.- Return type:
array- Parameters:
x (SparseTensor)
rho (float)