jaxDecomp: JAX Library for 3D Domain Decomposition and Parallel FFTs

Zenodo (CERN European Organization for Nuclear Research) 2026 Open access

Read full text
Authors: Wassim Kabalan, François Lanusse, Boucaud, Alexandre, É. Aubourg
DOI: 10.5281/zenodo.18512020
Cited by: 0

Abstract

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.

Fields of Study & Subjects

Parallel Computing and Optimization Techniques Computational Physics and Python Applications Matrix Theory and Algorithms

Publication details

Journal / source
Zenodo (CERN European Organization for Nuclear Research)
Publication date
7 February 2026