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
  • Index

Index

F | H | J | M | P | R | T

F

  • fftfreq3d() (in module jaxdecomp.fft)

H

  • halo_exchange() (in module jaxdecomp.halo)

J

  • jaxdecomp.fft
    • module
  • jaxdecomp.halo
    • module
  • jaxdecomp.transpose
    • module

M

  • module
    • jaxdecomp.fft
    • jaxdecomp.halo
    • jaxdecomp.transpose

P

  • pfft3d() (in module jaxdecomp.fft)
  • pifft3d() (in module jaxdecomp.fft)

R

  • rfftfreq3d() (in module jaxdecomp.fft)

T

  • transposeXtoY() (in module jaxdecomp.transpose)
  • transposeXtoZ() (in module jaxdecomp.transpose)
  • transposeYtoX() (in module jaxdecomp.transpose)
  • transposeYtoZ() (in module jaxdecomp.transpose)
  • transposeZtoX() (in module jaxdecomp.transpose)
  • transposeZtoY() (in module jaxdecomp.transpose)

© Copyright 2025, Wassim Kabalan, François Lanusse.

Built with Sphinx using a theme provided by Read the Docs.