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 or an explicit positional target.

mlx_lattice.ops.conv_transpose3d()

GenerativeConvTranspose3d

Generated transpose-convolution support or an explicit positional target.

mlx_lattice.ops.generative_conv_transpose3d()

NormalizedSubmConv3d

Input coordinate identity.

mlx_lattice.ops.normalized_subm_conv3d()

NormalizedConvTranspose3d

Transposed relation support or an explicit positional target.

mlx_lattice.ops.normalized_conv_transpose3d()

NormalizedGenerativeConvTranspose3d

Generated transpose-convolution support or an explicit positional target.

mlx_lattice.ops.normalized_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. The spatial axes are x, y, z with z varying fastest in flattened views. 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.NormalizedConvTranspose3d(in_channels, out_channels, *, kernel_size=2, stride=2, padding=0, dilation=1, bias=True, eps=1e-08)[source]

Bases: ConvTranspose3d

Weight-normalized sparse transpose convolution.

Parameters:
class mlx_lattice.nn.conv.NormalizedGenerativeConvTranspose3d(in_channels, out_channels, *, kernel_size=2, stride=2, bias=True, eps=1e-08)[source]

Bases: GenerativeConvTranspose3d

Weight-normalized generative sparse transpose convolution.

Parameters:
class mlx_lattice.nn.conv.NormalizedSubmConv3d(in_channels, out_channels, *, kernel_size=3, dilation=1, bias=True, eps=1e-08)[source]

Bases: SubmConv3d

Weight-normalized submanifold sparse convolution.

Parameters:
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)