Convolution modules

Convolution modules own dense floating weights and delegate execution to the functional sparse convolution API. Use them for model composition when weights are learned or stored as module parameters.

Module summary

Module

Coordinate support

Functional equivalent

Conv3d

Forward support or explicit target support.

mlx_lattice.ops.conv3d()

SubmConv3d

Input coordinate identity.

mlx_lattice.ops.subm_conv3d()

ConvTranspose3d

Transposed relation support.

mlx_lattice.ops.conv_transpose3d()

GenerativeConvTranspose3d

Generated transpose-convolution support.

mlx_lattice.ops.generative_conv_transpose3d()
class mlx_lattice.nn.conv.Conv3d(in_channels, out_channels, *, kernel_size=3, stride=1, padding=0, dilation=1, bias=True)[source]

Bases: Module

Sparse 3D convolution module.

The module owns a dense 5D weight with layout (C_out, Kx, Ky, Kz, C_in) and optional bias. Calling the module delegates to mlx_lattice.ops.conv3d(); coordinate support follows the same forward or explicit-target semantics as the functional API.

Parameters:

  • in_channels (int)

  • out_channels (int)

  • kernel_size (int | Sequence[int])

  • stride (int | Sequence[int])

  • padding (int | Sequence[int])

  • dilation (int | Sequence[int])

  • bias (bool)

to_quantized(group_size=None, bits=None, *, mode='affine', quantize_input=False)[source]
Return type:

QuantizedConv3d

Parameters:

  • group_size (int | None)

  • bits (int | None)

  • mode (str)

  • quantize_input (bool)

class mlx_lattice.nn.conv.ConvTranspose3d(in_channels, out_channels, *, kernel_size=2, stride=2, padding=0, dilation=1, bias=True)[source]

Bases: Module

Sparse 3D transpose-convolution module.

Calling the module delegates to mlx_lattice.ops.conv_transpose3d() and returns support generated by the transposed relation.

Parameters:

  • in_channels (int)

  • out_channels (int)

  • kernel_size (int | Sequence[int])

  • stride (int | Sequence[int])

  • padding (int | Sequence[int])

  • dilation (int | Sequence[int])

  • bias (bool)

to_quantized(group_size=None, bits=None, *, mode='affine', quantize_input=False)[source]
Return type:

QuantizedConvTranspose3d

Parameters:

  • group_size (int | None)

  • bits (int | None)

  • mode (str)

  • quantize_input (bool)

class mlx_lattice.nn.conv.GenerativeConvTranspose3d(in_channels, out_channels, *, kernel_size=2, stride=2, bias=True)[source]

Bases: Module

Generative sparse 3D transpose-convolution module.

Calling the module delegates to mlx_lattice.ops.generative_conv_transpose3d() and generates output coordinates from input rows and stride.

Parameters:

  • in_channels (int)

  • out_channels (int)

  • kernel_size (int | Sequence[int])

  • stride (int | Sequence[int])

  • bias (bool)

to_quantized(group_size=None, bits=None, *, mode='affine', quantize_input=False)[source]
Return type:

QuantizedGenerativeConvTranspose3d

Parameters:

  • group_size (int | None)

  • bits (int | None)

  • mode (str)

  • quantize_input (bool)

class mlx_lattice.nn.conv.SubmConv3d(in_channels, out_channels, *, kernel_size=3, dilation=1, bias=True)[source]

Bases: Module

Submanifold sparse 3D convolution module.

The output sparse tensor reuses input coordinate identity. Kernel sizes must be odd because the relation is centered on existing active rows.

Parameters:

  • in_channels (int)

  • out_channels (int)

  • kernel_size (int | Sequence[int])

  • dilation (int | Sequence[int])

  • bias (bool)

to_quantized(group_size=None, bits=None, *, mode='affine', quantize_input=False)[source]
Return type:

QuantizedSubmConv3d

Parameters:

  • group_size (int | None)

  • bits (int | None)

  • mode (str)

  • quantize_input (bool)