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 or target transposed	Generates transpose support or evaluates on explicit target coordinates.
generative_conv_transpose3d	generative	Generates support from input rows and stride.
normalized_subm_conv3d	submanifold	Reuses input support and applies weight-dependent normalization.
normalized_conv_transpose3d	transposed or target transposed	Applies weight-dependent normalization on generated or target support.
normalized_generative_conv_transpose3d	generative	Generates support and applies weight-dependent normalization.

Floating weights accept dense 5D layout (C_out, Kx, Ky, Kz, C_in). Spatial axes are x, y, z and flattened rows vary z fastest; this is the conv3d_o_xyz_i artifact layout. They also accept 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. The spatial axes are x, y, z; flattened kernel rows therefore have z varying fastest.

• 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, coordinates=None)

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])

• coordinates (SparseTensor | CoordinateMapKey | array | None)

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

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])

• coordinates (SparseTensor | CoordinateMapKey | array | None)

mlx_lattice.ops.conv.normalized_conv_transpose3d(x, weight, bias=None, *, kernel_size=2, stride=2, padding=0, dilation=1, eps=1e-08, coordinates=None)

Apply weight-normalized sparse transpose convolution.

Return type:

SparseTensor

Parameters:

• x (SparseTensor)

• weight (array)

• bias (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• padding (int | Sequence[int])

• dilation (int | Sequence[int])

• eps (float)

• coordinates (SparseTensor | CoordinateMapKey | array | None)

mlx_lattice.ops.conv.normalized_generative_conv_transpose3d(x, weight, bias=None, *, kernel_size=2, stride=2, eps=1e-08, coordinates=None)

Generate support with weight-normalized transpose convolution.

Return type:

SparseTensor

Parameters:

• x (SparseTensor)

• weight (array)

• bias (array | None)

• kernel_size (int | Sequence[int])

• stride (int | Sequence[int])

• eps (float)

• coordinates (SparseTensor | CoordinateMapKey | array | None)

mlx_lattice.ops.conv.normalized_subm_conv3d(x, weight, bias=None, *, kernel_size=3, dilation=1, eps=1e-08)

Apply weight-normalized convolution on input coordinate support.

For non-pointwise kernels the output is divided channel-wise by sqrt(conv(ones, weight**2) + eps) before bias is applied. This is the normalized sparse-convolution contract used by Gameleon’s attribute models. Pointwise kernels intentionally use the ordinary matrix product.

Return type:

SparseTensor

Parameters:

• x (SparseTensor)

• weight (array)

• bias (array | None)

• kernel_size (int | Sequence[int])

• dilation (int | Sequence[int])

• eps (float)

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])