torch-lattice

torch-lattice is the Torch/CUDA training-side companion to mlx-lattice. It keeps the sparse training and export workflow in PyTorch while producing MLIR artifact bundles that the MLX runtime can replay on Apple Silicon.

The library intentionally separates three concerns:

  • CUDA sparse operators and torch.nn modules for training and validation;

  • a small sparse tensor contract shared with the MLX side;

  • artifact tooling that lowers explicit Torch graphs to a portable lattice IR.

The primary data model matches MLX Lattice:

  • coordinates are integer rows ordered as (batch, x, y, z);

  • features are dense Torch tensors whose rows align one-to-one with coordinates;

  • sparse operators either preserve support, generate new support, or consume a caller-provided target support;

  • exported artifacts carry graph structure and named weights separately so the deployment runtime can rebuild the same sparse computation.