Convolution operations

Sparse convolution functions operate on mlx_lattice.SparseTensor objects and dense or packed weights. They build or reuse kernel relations and return a new sparse tensor whose coordinate support is determined by the operation:

Function	Relation kind	Output coordinates
conv3d	forward or target	Generated from input geometry, or taken from explicit target coordinates.
subm_conv3d	submanifold	Reuses input coordinate identity.
conv_transpose3d	transposed	Uses the transpose-convolution relation support.
generative_conv_transpose3d	generative	Generates support from input rows and stride.

Floating weights accept dense 5D layout (C_out, Kx, Ky, Kz, C_in) and mapped kernel-major layout (K, C_in, C_out). Packed quantized weights use QuantizedWeight.

Related pages

• Route predicates and backend details: Convolution routes

• Relation model: Coordinates and relations

• Sparse tensor output metadata: Sparse tensor model

• Module wrappers: Convolution modules

• Quantized weight storage: Quantized weights

mlx_lattice.ops.conv.conv3d(x, weight, bias=None, *, kernel_size=3, stride=1, padding=0, dilation=1, coordinates=None)

Apply sparse 3D convolution to a SparseTensor.

The pointwise 1x1x1 case with no explicit target coordinates uses a direct feature matrix multiply and preserves input coordinates. Other shapes build a sparse kernel relation and dispatch to the selected native CPU or Metal backend. coordinates may provide an explicit target lattice for target-aligned convolution.

Parameters:

• x (SparseTensor) – Input sparse tensor with feature shape (N_in, C_in).

• weight (array | QuantizedWeight) – Floating weight or packed QuantizedWeight. Floating weights accept (C_out, C_in), (K, C_in, C_out), or (C_out, Kx, Ky, Kz, C_in) depending on kernel geometry.

• bias (array | None) – Optional (C_out,) bias matching output feature dtype.

• kernel_size (int | Sequence[int]) – 3D kernel size.

• stride (int | Sequence[int]) – 3D convolution stride. Output sparse stride is x.stride * stride.

• padding (int | Sequence[int]) – 3D padding used by the relation builder.

• dilation (int | Sequence[int]) – 3D dilation used by the relation builder.

• coordinates (SparseTensor | CoordinateMapKey | array | None) – Optional explicit target support as a sparse tensor, coordinate key, or coordinate array.

Return type:

SparseTensor

Returns:

Sparse tensor with output features (N_out, C_out). If coordinates is supplied, output coordinates are exactly that target support; otherwise they are generated by the forward relation.

mlx_lattice.ops.conv.conv_transpose3d(x, weight, bias=None, *, kernel_size=2, stride=2, padding=0, dilation=1)

Apply sparse transpose convolution using a transposed kernel relation.

The output sparse stride is x.stride / stride and each component of stride must divide the corresponding input stride. Output support is produced by the transposed relation builder.

Return type:

SparseTensor

Parameters:

• x (SparseTensor)

• weight (array | QuantizedWeight)

• bias (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])

mlx_lattice.ops.conv.generative_conv_transpose3d(x, weight, bias=None, *, kernel_size=2, stride=2)

Generate output coordinates from a sparse transpose-convolution rule.

This operation expands support directly from input rows and transpose stride. It is useful for sparse decoders that need generated coordinates before subsequent target-aligned operations.

Return type:

SparseTensor

Parameters:

• x (SparseTensor)

• weight (array | QuantizedWeight)

• bias (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

mlx_lattice.ops.conv.subm_conv3d(x, weight, bias=None, *, kernel_size=3, dilation=1)

Apply submanifold sparse convolution without changing coordinates.

Submanifold convolution uses an odd kernel, stride 1, and no padding expansion. The result reuses input coordinate identity, so downstream feature-only operations and relation caches observe the same support.

Return type:

SparseTensor

Parameters:

• x (SparseTensor)

• weight (array | QuantizedWeight)

• bias (array | None)

• kernel_size (int | Sequence[int])

• dilation (int | Sequence[int])