JaxDecomp Documentation

A high-performance, differentiable distributed FFT library for JAX.

User Guide

User Guide

• Installation
  • 1. Pure JAX Version (Easy / Recommended)
  • 2. cuDecomp Backend (Advanced / HPC)
  • Machine-Specific Installation Notes
  • Backend Selection at Runtime
  • cuDecomp Transpose Communication Backends
  • Notes on Performance
• Basic Usage
• Understanding Domain Decomposition in jaxDecomp
  • What is pdims?
  • Slab vs Pencil: Tradeoffs
  • Slab Decomposition
  • Pencil Decomposition
  • Dynamically Generating pdims
  • Recommendation
  • Creating the JAX Mesh and Sharding
  • TL;DR
• Distributed 3D FFTs in jaxDecomp
  • Overview of the FFT Algorithm
  • Transpositions Between Axes
  • Pencil Decomposition Strategy
  • Slab Decomposition Strategy
  • Non-Contiguous Global Transpositions
  • Summary
• Halo Exchange in jaxDecomp
  • Halo Exchange Process
  • Boundary Conditions
  • Padding and Halo Exchange
  • Implementation Details
  • Summary
• Transpositions in jaxDecomp
  • What is a Global Transpose?
  • Visual Illustration
  • Contiguous vs Non-Contiguous Transpositions
  • API Example
  • Summary
• Caveats and Workarounds: Autodiff + SPMD Sharding with jaxDecomp
  • 1. Background
  • 2. jacfwd with Parallel FFT
  • 3. jacrev with Parallel FFT
  • 4. grad of a Scalar-Reduced FFT
  • 5. Summary of Best Practices
  • 6. Conclusion
• XLA Sharding Configuration Guide
  • Shardy Partitioner Configuration
  • Sharding Spec Must Match Mesh
  • Auto vs Explicit Axis Types

API Reference

API Reference

• API Reference
  • FFT
  • Halo Exchange
  • Transpositions