Quantized convolution modules

Quantized convolution modules store packed affine int4/int8 weights and keep floating-point sparse activations. Coordinate semantics match the floating module with the same geometry.

Module summary

Module

Floating source

Coordinate support

QuantizedConv3d

Conv3d

Forward or explicit target support.

QuantizedSubmConv3d

SubmConv3d

Input coordinate identity.

QuantizedConvTranspose3d

ConvTranspose3d

Transposed relation support.

QuantizedGenerativeConvTranspose3d

GenerativeConvTranspose3d

Generated transpose-convolution support.
class mlx_lattice.nn.quantized_conv.QuantizedConv3d(in_channels, out_channels, *, kernel_size=3, stride=1, padding=0, dilation=1, bias=True, group_size=None, bits=4)[source]

Bases: _QuantizedConvBase

Affine weight-quantized sparse 3D convolution module.

Weights are stored as packed int4/int8 affine QuantizedWeight metadata. Activations remain floating point. Coordinate semantics match mlx_lattice.nn.Conv3d.

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)

  • group_size (int | None)

  • bits (int)

classmethod from_conv(source, group_size=None, bits=4)[source]
Return type:

QuantizedConv3d

Parameters:
class mlx_lattice.nn.quantized_conv.QuantizedConvTranspose3d(in_channels, out_channels, *, kernel_size=2, stride=2, padding=0, dilation=1, bias=True, group_size=None, bits=4)[source]

Bases: _QuantizedConvBase

Affine weight-quantized sparse transpose-convolution module.

Activations remain floating point and weight storage is packed affine int4/int8. Coordinate generation matches ConvTranspose3d.

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)

  • group_size (int | None)

  • bits (int)

classmethod from_conv(source, group_size=None, bits=4)[source]
Return type:

QuantizedConvTranspose3d

Parameters:
class mlx_lattice.nn.quantized_conv.QuantizedGenerativeConvTranspose3d(in_channels, out_channels, *, kernel_size=2, stride=2, bias=True, group_size=None, bits=4)[source]

Bases: _QuantizedConvBase

Affine weight-quantized generative transpose-convolution module.

The module stores packed affine weights and delegates coordinate generation to the generative transpose-convolution relation.

Parameters:

  • in_channels (int)

  • out_channels (int)

  • kernel_size (int | Sequence[int])

  • stride (int | Sequence[int])

  • bias (bool)

  • group_size (int | None)

  • bits (int)

classmethod from_conv(source, group_size=None, bits=4)[source]
Return type:

QuantizedGenerativeConvTranspose3d

Parameters:
class mlx_lattice.nn.quantized_conv.QuantizedSubmConv3d(in_channels, out_channels, *, kernel_size=3, dilation=1, bias=True, group_size=None, bits=4)[source]

Bases: _QuantizedConvBase

Affine weight-quantized submanifold convolution module.

Coordinate identity is preserved exactly as in SubmConv3d.

Parameters:

  • in_channels (int)

  • out_channels (int)

  • kernel_size (int | Sequence[int])

  • dilation (int | Sequence[int])

  • bias (bool)

  • group_size (int | None)

  • bits (int)

classmethod from_conv(source, group_size=None, bits=4)[source]
Return type:

QuantizedSubmConv3d

Parameters: