Coordinate utilities

Coordinate utilities cover set operations, ordering, occupancy, and alignment. For sparse tensor identity semantics, read Sparse tensor model. For value-aligned sparse algebra, read Coordinate-aligned sparse algebra and the operation API at Sparse tensor algebra.

Set operations

class mlx_lattice.core.coords.set_ops.CoordinateSet(coords, active_rows)

Bases: object

Capacity coordinate buffer with a lazy active count.

Parameters:

• coords (array)

• active_rows (array)

coords: array

active_rows: array

property capacity: int

property active_count: array

mlx_lattice.core.coords.set_ops.downsample_coords(coords, stride=2)

Downsample coordinates by integer spatial stride.

Return type:

CoordinateSet

Parameters:

• coords (array)

• stride (int | Sequence[int])

mlx_lattice.core.coords.set_ops.union_coords(lhs, rhs)

Return the coordinate-set union of two coordinate arrays.

Return type:

CoordinateSet

Parameters:

• lhs (array)

• rhs (array)

mlx_lattice.core.coords.set_ops.intersection_coords(lhs, rhs)

Return coordinates present in both input coordinate arrays.

Return type:

CoordinateSet

Parameters:

• lhs (array)

• rhs (array)

mlx_lattice.core.coords.set_ops.lookup_coords(coords, queries)

Map query coordinates to row indices in coords.

Missing query rows are encoded as -1.

Return type:

array

Parameters:

• coords (array)

• queries (array)

mlx_lattice.core.coords.set_ops.contains_coords(coords, queries)

Return a boolean mask indicating which queries exist in coords.

Return type:

array

Parameters:

• coords (array)

• queries (array)

mlx_lattice.core.coords.set_ops.inverse_map(source, target)

Return row indices that gather target rows from source.

Return type:

array

Parameters:

• source (array)

• target (array)

Ordering

class mlx_lattice.core.coords.ordering.CoordinateOrdering(coords, order, inverse_rows)

Bases: object

Sorted coordinate rows and the row permutation used to produce them.

Parameters:

• coords (array)

• order (array)

• inverse_rows (array)

coords: array

order: array

inverse_rows: array

mlx_lattice.core.coords.ordering.morton_codes(coords)

Return Morton/Z-order codes for batched sparse coordinates.

Return type:

array

Parameters:

coords (array)

mlx_lattice.core.coords.ordering.morton_order(coords)

Return row indices that sort coordinates by Morton order.

Return type:

array

Parameters:

coords (array)

mlx_lattice.core.coords.ordering.morton_sort_coords(coords)

Return coordinates sorted by Morton order and the ordering indices.

Return type:

CoordinateOrdering

Parameters:

coords (array)

Occupancy

class mlx_lattice.core.coords.occupancy.SparseOccupancy(coords, active_rows, occupancy)

Bases: object

Downsampled coordinates plus an 8-bit child occupancy mask per row.

Parameters:

• coords (array)

• active_rows (array)

• occupancy (array)

coords: array

active_rows: array

occupancy: array

property capacity: int

property active_count: array

class mlx_lattice.core.coords.occupancy.OccupancyExpansion(coords, active_rows, parent_rows, child_indices)

Bases: object

Expanded child coordinates with parent row and child-index metadata.

Parameters:

• coords (array)

• active_rows (array)

• parent_rows (array)

• child_indices (array)

coords: array

active_rows: array

parent_rows: array

child_indices: array

property capacity: int

property active_count: array

mlx_lattice.core.coords.occupancy.occupancy_downsample(coords, active_rows=None)

Downsample coordinates and record occupied children for each parent.

Return type:

SparseOccupancy

Parameters:

• coords (array)

• active_rows (array | None)

mlx_lattice.core.coords.occupancy.occupancy_expand(coords, occupancy, active_rows=None)

Expand occupied child coordinates from a parent occupancy mask.

Return type:

OccupancyExpansion

Parameters:

• coords (array)

• occupancy (array)

• active_rows (array | None)

mlx_lattice.core.coords.occupancy.child_coords_from_indices(parent_coords, child_indices)

Compute child coordinates from parent coordinates and child indices.

Return type:

array

Parameters:

• parent_coords (array)

• child_indices (array)

Alignment

type mlx_lattice.core.coords.alignment.SparseJoin = Literal['inner', 'left', 'right', 'outer']

class mlx_lattice.core.coords.alignment.SparseAlignment(coords, active_rows, lhs_rows, rhs_rows)

Bases: object

Coordinate value alignment for two sparse tensors.

coords is the joined coordinate support. lhs_rows and rhs_rows gather features from the left and right tensors into that support; -1 marks a missing value for the selected join mode.

Parameters:

• coords (array)

• active_rows (array)

• lhs_rows (array)

• rhs_rows (array)

coords: array

active_rows: array

lhs_rows: array

rhs_rows: array

property capacity: int

property active_count: array

mlx_lattice.core.coords.alignment.build_sparse_alignment(lhs_coords, lhs_active_rows, rhs_coords, rhs_active_rows, *, join='inner')

Build value-aligned row maps between two coordinate arrays.

Coordinates must have shape (N, 4) and matching dtype. Active-row scalars describe the valid prefix of each coordinate buffer. The native builder returns joined coordinates and row maps for inner, left, right, or outer support.

Return type:

SparseAlignment

Parameters:

• lhs_coords (array)

• lhs_active_rows (array)

• rhs_coords (array)

• rhs_active_rows (array)

• join (SparseJoin)