jaxDecomp issues

Fix 3D FFT

2

This PR will aim to fix the issues with 3D FFTs, right now, it just changes the FFT test to test the result of the forward fft, and this test...

EiffL

Implement Transpose primitives

2

First pull request for Tranpose ops They work for cubes for now, (Pencils and slabs)

ASKabalan

Fix slice unpad

3

Fixed slice_unpad, now it supports complex numbers

ASKabalan

Things to update for version 0.0.1

1

- [ ] slide_unpad does not seem to be too happy with getting complex numbers as inputs: ``` File "/mnt/home/flanusse/repo/jaxDecomp/scripts/demo.py", line 58, in recarray = slice_unpad(exchanged_reduced, padding_width, pdims) jaxlib.xla_extension.XlaRuntimeError: INTERNAL:...

EiffL

Easier way to chose cuda versions for end users

4

Currently in the CMake, the cuda version is set to be 12.2. jaxDecomp (and cuDecomp) can be compiled with 11.8 (no specific cuda 12 code) JAX 0.4.26 and above no...

ASKabalan

enhancement

FFTs are not working properly

12

Comparing the 3D FFT computed by jaxdecomp and manually in jax, I realized that the result of fft3d does not match with the non-distributed version. This could be due to...

EiffL

bug

Currently unsupported features and caveats (aka the big TODO list)

1

Because implementing all features is not necessarily interesting unless there is a need for it, here are the current restrictions of the code in its current version. **All these restrictions...

EiffL