Relations

Relation objects connect coordinate semantics to backend execution. Start with Coordinates and relations for the conceptual model, then use this page for exact class members and builders. Convolution and pooling route details are documented in Convolution routes and Pooling routes.

Stable application code should normally depend on KernelSpec, KernelRelation, NeighborRelation, and the functional relation builders. Low-level CSR views, implicit-GEMM views, sorted implicit-GEMM views, and relation contracts are exposed for inspection and debugging, but their exact layout is provisional backend metadata.

Relation specifications

class mlx_lattice.core.relations.specs.KernelSpec(size=3, stride=1, padding=0, dilation=1)

Bases: object

Normalized 3D kernel geometry for sparse relation builders.

KernelSpec accepts integers or three-item sequences for size, stride, padding, and dilation and stores all values as triples. The object is hashable, so coordinate managers can use it directly as part of relation cache keys.

Parameters:

• size (tuple[int, int, int])

• stride (tuple[int, int, int])

• padding (tuple[int, int, int])

• dilation (tuple[int, int, int])

size: tuple[int, int, int]

stride: tuple[int, int, int]

padding: tuple[int, int, int]

dilation: tuple[int, int, int]

property volume: int

Number of offsets in the dense kernel footprint.

property is_pointwise: bool

Whether this spec is an identity 1x1x1 pointwise mapping.

property is_centered_submanifold: bool

Whether this spec is an odd, stride-1 submanifold footprint.

Relation views

The view classes below are advanced/provisional. They are useful when diagnosing relation construction or backend route selection; they should not be used as a long-term application storage format.

class mlx_lattice.core.relations.views.RelationEdges(in_rows, out_rows, kernel_ids)

Bases: object

Diagnostic edge arrays for a logical sparse kernel relation.

Each edge stores (in_row, out_row, kernel_id) at the same index in the three arrays. Edge arrays are useful for debugging and for generic sparse traversal routes; CSR views provide grouped execution order.

Parameters:

• in_rows (array)

• out_rows (array)

• kernel_ids (array)

in_rows: array

out_rows: array

kernel_ids: array

property capacity: int

class mlx_lattice.core.relations.views.RelationCSRView(row_offsets, edge_ids=None)

Bases: object

CSR execution view over relation edge arrays.

row_offsets has shape (rows + 1,) and int32 dtype. Optional edge_ids remaps CSR order back to the canonical edge arrays.

Parameters:

• row_offsets (array)

• edge_ids (array | None)

row_offsets: array

edge_ids: array | None

mlx_lattice.core.relations.views.RelationView

alias of RelationCSRView

class mlx_lattice.core.relations.views.RelationImplicitGemmView(out_in_map, row_masks)

Bases: object

Dense output-row by kernel-offset map for implicit-GEMM execution.

out_in_map[o, k] stores the input row that contributes to output row o at kernel offset k. Missing contributors are represented by the native builder’s sentinel value. row_masks stores bit masks describing populated kernel offsets for each output row.

Parameters:

• out_in_map (array)

• row_masks (array)

out_in_map: array

row_masks: array

class mlx_lattice.core.relations.views.RelationSortedImplicitGemmView(sorted_out_in_map, sorted_kv_out_in_map, reorder_rows, tile_masks)

Bases: object

Tile-sorted implicit-GEMM execution view for TensorOps kernels.

Parameters:

• sorted_out_in_map (array)

• sorted_kv_out_in_map (array)

• reorder_rows (array)

• tile_masks (array)

sorted_out_in_map: array

sorted_kv_out_in_map: array

reorder_rows: array

tile_masks: array

class mlx_lattice.core.relations.views.SparseRelationContract(counts, kernel_offsets, out_coords=None, n_in_capacity=None, n_out_capacity=None, n_kernels=None, source_coords=None, source_active_rows=None, target_coords=None, target_active_rows=None, stride=(1, 1, 1), padding=(0, 0, 0), kind='forward')

Bases: object

Logical sparse relation contract shared by all execution views.

The contract records counts, capacities, optional source/target coordinate buffers, kernel offsets, stride, padding, and relation kind. Backend routes use this metadata to validate specialized kernels before consuming CSR or implicit-GEMM views.

Parameters:

• counts (array)

• kernel_offsets (tuple[tuple[int, int, int], ...])

• out_coords (array | None)

• n_in_capacity (int | None)

• n_out_capacity (int | None)

• n_kernels (int | None)

• source_coords (array | None)

• source_active_rows (array | None)

• target_coords (array | None)

• target_active_rows (array | None)

• stride (tuple[int, int, int])

• padding (tuple[int, int, int])

• kind (Literal['forward', 'target', 'submanifold', 'transposed', 'generative'])

counts: array

kernel_offsets: tuple[tuple[int, int, int], ...]

out_coords: array | None

n_in_capacity: int | None

n_out_capacity: int | None

n_kernels: int | None

source_coords: array | None

source_active_rows: array | None

target_coords: array | None

target_active_rows: array | None

stride: tuple[int, int, int]

padding: tuple[int, int, int]

kind: Literal['forward', 'target', 'submanifold', 'transposed', 'generative']

property edge_count: array

property out_count: array

class mlx_lattice.core.relations.views.NeighborEdges(query_rows, source_rows, neighbor_ids)

Bases: object

Semantic neighbor edge arrays for query/source relations.

Each edge stores the query row, source row, and neighbor rank or identifier for geometric neighbor queries such as kNN and radius search.

Parameters:

• query_rows (array)

• source_rows (array)

• neighbor_ids (array)

query_rows: array

source_rows: array

neighbor_ids: array

property capacity: int

class mlx_lattice.core.relations.views.KernelRelation(in_rows, out_rows, kernel_ids, *, row_offsets=None, counts=None, in_row_offsets=None, in_edge_ids=None, kernel_row_offsets=None, kernel_edge_ids=None, kernel_offsets=(), out_coords=None, n_in_capacity=None, n_out_capacity=None, n_kernels=None, source_coords=None, source_active_rows=None, target_coords=None, target_active_rows=None, stride=(1, 1, 1), padding=(0, 0, 0), kind='forward', implicit_gemm=None, sorted_implicit_gemm=None)

Bases: object

Sparse kernel relation with multiple execution-oriented views.

A relation stores semantic edges (in_row, out_row, kernel_id) plus CSR views grouped by output row, input row, and kernel id. Forward convolution, pooling, and backend-specific fast paths consume these views.

Parameters:

• in_rows (mx.array)

• out_rows (mx.array)

• kernel_ids (mx.array)

• row_offsets (mx.array | None)

• counts (mx.array | None)

• in_row_offsets (mx.array | None)

• in_edge_ids (mx.array | None)

• kernel_row_offsets (mx.array | None)

• kernel_edge_ids (mx.array | None)

• kernel_offsets (tuple[Triple, ...])

• out_coords (mx.array | None)

• n_in_capacity (int | None)

• n_out_capacity (int | None)

• n_kernels (int | None)

• source_coords (mx.array | None)

• source_active_rows (mx.array | None)

• target_coords (mx.array | None)

• target_active_rows (mx.array | None)

• stride (Triple)

• padding (Triple)

• kind (RelationKind)

• implicit_gemm (RelationImplicitGemmView | None)

• sorted_implicit_gemm (RelationSortedImplicitGemmView | None)

contract: SparseRelationContract

edges: RelationEdges

output_csr: RelationCSRView

input_csr: RelationCSRView

kernel_csr: RelationCSRView

implicit_gemm: RelationImplicitGemmView | None

sorted_implicit_gemm: RelationSortedImplicitGemmView | None

property edge_capacity: int

Static edge-buffer capacity.

property counts: array

Native counts array [edge_count, out_count].

property edge_count: array

property out_count: array

property kernel_offsets: tuple[tuple[int, int, int], ...]

property out_coords: array | None

property n_in_capacity: int | None

property n_out_capacity: int | None

property n_kernels: int | None

property row_offsets: array

property in_row_offsets: array

property in_edge_ids: array

property kernel_row_offsets: array

property kernel_edge_ids: array

property out_view: RelationCSRView

property in_view: RelationCSRView

property kernel_view: RelationCSRView

require_implicit_gemm()

Return or lazily build the implicit-GEMM execution view.

Return type:

RelationImplicitGemmView

require_sorted_implicit_gemm()

Return or lazily build the tile-sorted implicit-GEMM view.

Return type:

RelationSortedImplicitGemmView

class mlx_lattice.core.relations.views.NeighborRelation(query_rows, source_rows, neighbor_ids, distances, *, row_offsets=None, counts=None, n_query_capacity=None, n_source_capacity=None, max_neighbors=None)

Bases: object

Sparse query/source neighbor relation with row-offset metadata.

Neighbor relations are not convolution kernel relations. They describe geometric connectivity between query rows and source rows, carry distances in edge order, and expose row offsets grouped by query row.

Parameters:

• query_rows (mx.array)

• source_rows (mx.array)

• neighbor_ids (mx.array)

• distances (array)

• row_offsets (array)

• counts (array)

• n_query_capacity (int | None)

• n_source_capacity (int | None)

• max_neighbors (int | None)

edges: NeighborEdges

distances: array

row_offsets: array

counts: array

n_query_capacity: int | None

n_source_capacity: int | None

max_neighbors: int | None

property edge_capacity: int

property edge_count: array

property query_count: array

Relation builders

type mlx_lattice.core.coords.builders.NativeKernelRelation = tuple[array, array, array, array, array, array, array, array, array, array]

type mlx_lattice.core.coords.builders.NativeNeighborRelation = tuple[array, array, array, array, array, array]

type mlx_lattice.core.coords.builders.NativeImplicitGemmView = tuple[array, array]

mlx_lattice.core.coords.builders.build_relation_implicit_gemm_view(relation)

Build the dense output-to-input map used by implicit-GEMM kernels.

Return type:

RelationImplicitGemmView

Parameters:

relation (KernelRelation)

mlx_lattice.core.coords.builders.build_target_kernel_relation(coords, target_coords, *, active_rows=None, target_active_rows=None, kernel_size=3, stride=1, padding=0, dilation=1)

Build a sparse kernel relation from source coords to target coords.

Return type:

KernelRelation

Parameters:

• coords (array)

• target_coords (array)

• active_rows (array | None)

• target_active_rows (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])

mlx_lattice.core.coords.builders.kernel_offsets(kernel_size, dilation=1)

Enumerate spatial offsets for a dense 3D kernel footprint.

Return type:

tuple[tuple[int, int, int], ...]

Parameters:

• kernel_size (int | Sequence[int])

• dilation (int | Sequence[int])

mlx_lattice.core.coords.builders.build_kernel_relation(coords, active_rows=None, kernel_size=3, stride=1, padding=0, dilation=1)

Build a forward sparse convolution/pooling relation.

Return type:

KernelRelation

Parameters:

• coords (array)

• active_rows (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])

mlx_lattice.core.coords.builders.build_submanifold_kernel_relation(coords, active_rows=None, kernel_size=3, dilation=1)

Build a submanifold relation whose output support is coords.

Return type:

KernelRelation

Parameters:

• coords (array)

• active_rows (array | None)

• kernel_size (int | Sequence[int])

• dilation (int | Sequence[int])

mlx_lattice.core.coords.builders.build_generative_relation(coords, active_rows=None, kernel_size=2, stride=2)

Build a generative transpose-convolution relation.

Return type:

KernelRelation

Parameters:

• coords (array)

• active_rows (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

mlx_lattice.core.coords.builders.build_transposed_kernel_relation(coords, active_rows=None, kernel_size=2, stride=2, padding=0, dilation=1)

Build a sparse transpose-convolution relation.

Return type:

KernelRelation

Parameters:

• coords (array)

• active_rows (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])

mlx_lattice.core.coords.builders.build_knn_relation(source_coords, query_coords=None, *, source_active_rows=None, query_active_rows=None, k)

Build a k-nearest-neighbor relation between source and query coords.

Return type:

NeighborRelation

Parameters:

• source_coords (array)

• query_coords (array | None)

• source_active_rows (array | None)

• query_active_rows (array | None)

• k (int)

mlx_lattice.core.coords.builders.build_radius_relation(source_coords, query_coords=None, *, source_active_rows=None, query_active_rows=None, radius, max_neighbors=None)

Build a radius-neighborhood relation between source and query coords.

Return type:

NeighborRelation

Parameters:

• source_coords (array)

• query_coords (array | None)

• source_active_rows (array | None)

• query_active_rows (array | None)

• radius (float)

• max_neighbors (int | None)