Published September 26, 2026
| Version v0.3.4
Software
Open
jaxDecomp: JAX Library for 3D Domain Decomposition and Parallel FFTs
Authors/Creators
Description
JAX reimplementation and bindings for NVIDIA's cuDecomp library, enabling multi-node parallel FFTs and halo exchanges directly in low-level NCCL/CUDA-Aware MPI from JAX code. jaxDecomp provides both a pure JAX backend for easy deployment and an optional cuDecomp backend for advanced MPI/NCCL communication patterns. All functions are JIT-compatible and support automatic differentiation.
Notes
If you use this software, please cite it as below.
Files
DifferentiableUniverseInitiative/jaxDecomp-v0.3.4.zip
Files
(43.3 MB)
| Name | Size | Download all |
|---|---|---|
|
md5:1b7ac943ce7ceb2b661b6c23f4fa8943
|
43.3 MB | Preview Download |
Additional details
Related works
- Is supplement to
- Software: https://github.com/DifferentiableUniverseInitiative/jaxDecomp/tree/v0.3.4 (URL)