JaxDecomp
User Guide
Installation
Basic Usage
Understanding Domain Decomposition in
jaxDecomp
Distributed 3D FFTs in
jaxDecomp
Halo Exchange in
jaxDecomp
Transpositions in
jaxDecomp
Caveats and Workarounds: Autodiff + SPMD Sharding with
jaxDecomp
XLA Sharding Configuration Guide
API Reference
API Reference
JaxDecomp
Overview: module code
All modules for which code is available
jaxdecomp.fft
jaxdecomp.halo
jaxdecomp.transpose