Coordinate management

Coordinate managers own coordinate arrays, active-row scalars, and relation caches. A CoordinateMapKey is only valid inside the manager that created it. This lets operations compare coordinate identity without comparing every row.

The cache key for a kernel relation includes:

• input coordinate key;

• optional explicit target coordinate key;

• normalized kernel geometry;

• relation kind, such as forward, target, submanifold, transposed, or generative.

This page documents the canonical manager API. Higher-level users usually construct managers indirectly by creating mlx_lattice.SparseTensor objects.

Related pages

• Sparse tensor identity: Sparse tensor model

• Relation views: Coordinates and relations

• Relation classes: Relations

• Functional relation builders: Relation operations

class mlx_lattice.core.coords.manager.CoordinateMapKey(id, stride, manager_id)

Bases: object

Opaque coordinate identity key owned by a CoordinateManager.

Keys are lightweight metadata objects. A key is valid only for the manager whose manager_id it stores. The stride field is part of coordinate identity because the same integer rows represent different lattice cells at different spatial strides.

Parameters:

• id (int)

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

• manager_id (int)

id: int

stride: tuple[int, int, int]

manager_id: int

class mlx_lattice.core.coords.manager.CoordinateManager(_next_id=0, _coords=<factory>, _identity_keys=<factory>, _default_identity_keys=<factory>, _active_rows=<factory>, _kernel_relations=<factory>)

Bases: object

Owns coordinate arrays and caches sparse relations by identity.

Managers are normally created by SparseTensor. Reusing a manager and key lets convolutions, pooling, and aligned sparse algebra reuse native relation metadata instead of rebuilding it for every call. Cache entries are keyed by coordinate identity, optional target identity, kernel geometry, and relation kind.

Parameters:

• _next_id (int)

• _coords (dict[CoordinateMapKey, array])

• _identity_keys (dict[tuple[int, int, tuple[int, int, int]], CoordinateMapKey])

• _default_identity_keys (dict[tuple[int, tuple[int, int, int]], CoordinateMapKey])

• _active_rows (dict[CoordinateMapKey, array])

• _kernel_relations (dict[tuple[CoordinateMapKey, CoordinateMapKey | None, KernelSpec, str], KernelRelation])

insert_coords(coords, stride=1, active_rows=None)

Register a coordinate array by object identity and stride.

Parameters:

• coords (array) – Integer coordinates with shape (N, 4).

• stride (int | Sequence[int]) – Spatial lattice stride for the coordinate rows.

• active_rows (array | None) – Optional int32 scalar (1,) active-row count.

Return type:

CoordinateMapKey

Returns:

A manager-owned key. Re-inserting the same coordinate object, active-row object, and stride returns the existing key.

owns(key)

Return whether key belongs to this manager.

Return type:

bool

Parameters:

key (CoordinateMapKey)

coords(key)

Return the coordinate array registered for key.

Return type:

array

Parameters:

key (CoordinateMapKey)

active_rows(key)

Return the active-row scalar registered for key.

Return type:

array

Parameters:

key (CoordinateMapKey)

inverse_map(source, target)

Return row indices that gather target coordinates from source.

Missing target rows are encoded by the native set-operation contract. Both keys must belong to this manager.

Return type:

array

Parameters:

• source (CoordinateMapKey)

• target (CoordinateMapKey)

kernel_relation(key, *, kernel_size=3, stride=1, padding=0, dilation=1)

Build or reuse a forward sparse kernel relation.

The relation maps input rows to output rows generated by applying kernel_size, stride, padding, and dilation to the input coordinate support. The resulting relation is cached for the exact normalized geometry.

Return type:

KernelRelation

Parameters:

• key (CoordinateMapKey)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])

generative_relation(key, *, kernel_size=2, stride=2)

Build or reuse a generative transpose-convolution relation.

Generative relations create output support from each input coordinate and transpose-convolution stride without requiring an existing target coordinate set.

Return type:

KernelRelation

Parameters:

• key (CoordinateMapKey)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

submanifold_kernel_relation(key, *, kernel_size=3, dilation=1)

Build or reuse a submanifold sparse kernel relation.

The relation fixes output support to the input coordinate identity. It is the relation semantic used by submanifold convolution rather than a generic forward relation with output coordinates discarded.

Return type:

KernelRelation

Parameters:

• key (CoordinateMapKey)

• kernel_size (int | Sequence[int])

• dilation (int | Sequence[int])

transposed_kernel_relation(key, *, kernel_size=2, stride=2, padding=0, dilation=1)

Build or reuse a transpose-convolution relation.

The relation represents the sparse transpose of the forward convolution geometry and carries output coordinates plus CSR views for native execution.

Return type:

KernelRelation

Parameters:

• key (CoordinateMapKey)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])

target_kernel_relation(key, target_key, *, kernel_size=3, stride=1, padding=0, dilation=1)

Build or reuse a relation from one key to explicit target coords.

target_key fixes the output support. Only input rows that connect to those target rows through the kernel geometry contribute edges.

Return type:

KernelRelation

Parameters:

• key (CoordinateMapKey)

• target_key (CoordinateMapKey)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])