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.

See ops/README.md for the math motivation behind the FFT-based ops, and docs-tracker.md at the repo root for the documentation coverage plan.

• 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