API Reference

Organised bottom-up: low-level convolution primitives first, then the mixer modules that compose them, then full networks, then the parallel, core utility, and experiments layers.

Start with the Ops Primer for the math motivation behind the FFT-based ops (the convolution theorem, the linear/circular flavours, and a decision tree for picking a function). docs-tracker.md at the repo root tracks the documentation coverage plan.

• Ops Primer
  • Why FFT convolution?
  • File map
  • Naming convention
  • Shape conventions
  • Choosing a function: a decision tree
  • Numerical notes
  • Related modules
• Ops
  • FFT convolutions (reference fp32)
  • FFT convolutions (CUDA-accelerated)
  • Direct 1D causal convolutions (CUDA-accelerated)
  • Circular FFT convolutions
  • Chunking utilities
  • Mixed boundary-condition FFT convolutions
• Modules
  • Mixers
  • Convolutions
  • Kernels & filters
  • Normalization
  • Position encoding & patching
  • Gating & conditioning
  • Residual blocks
  • Schedulers
• Networks
  • Classification & general-purpose
  • Baselines
• Parallel
  • init_parallel_state
  • zigzag_split_across_group_ranks
  • zigzag_gather_from_group_ranks
  • setup_rank0_logging
  • AllToAllSingleFunction
• Core
  • Lazy configuration
  • Initialisation helpers
  • QK normalization & rotary position embeddings
  • QuACK capability probe
  • Testing helpers
• Experiments
  • Entry points
  • Configuration dataclasses
  • Lightning wrappers
  • Callbacks
  • Data modules
  • Utilities