Neural network modules¶

mlx_lattice.nn mirrors the functional operation API with mlx.nn.Module-style parameter ownership. Modules accept and return mlx_lattice.SparseTensor for sparse operations, except global pooling modules, which return dense (B, C) MLX arrays.

Use this section when you need module constructors, parameter names, to_quantized behavior, or the exact relationship between module wrappers and functional operations.

Module map¶

Feature set	Module API	Functional/reference pages
Sparse convolution	Convolution modules	Convolution operations, Convolution routes
Quantized sparse convolution	Quantized convolution modules	Quantized weights, Quantization routes
Sparse feature modules	Feature modules	Feature operations, Sparse tensor model
Quantized sparse feature modules	Quantized feature modules	Feature operations, Quantized weights
Sparse pooling	Pooling modules	Pooling operations, Pooling routes

Coordinate behavior¶

Module family	Coordinate effect	Output type
`Conv3d`	Generates forward or explicit target support.	`SparseTensor`
`SubmConv3d`	Preserves input coordinate identity.	`SparseTensor`
Transpose convolution modules	Generate expanded sparse support.	`SparseTensor`
Feature modules	Preserve coordinate identity.	`SparseTensor`
Local pooling modules	Generate pooled support from a kernel relation.	`SparseTensor`
Global pooling modules	Reduce by `batch_counts`.	Dense MLX array