Skip to content

Make jax_cuda kernels deterministic and allow using multiple devices with the FFI kernel - #405

Open
ASKabalan wants to merge 5 commits into
mainfrom
ASKabalan
Open

ASKabalan wants to merge 5 commits into
mainfrom
ASKabalan

Conversation

@ASKabalan

Copy link
Copy Markdown
Collaborator

Fixing the multi device (DeviceOrd) bug where all devices are assigned the same cudaDevice
Reworking the XLA parameter aliasing to ensure determinism for gradient based algorithms

This closes #389

@codecov

codecov Bot commented Sep 13, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.10%. Comparing base (86a4dff) to head (a47f232).

Additional details and impacted files
@@           Coverage Diff           @@
##             main     #405   +/-   ##
=======================================
  Coverage   96.10%   96.10%           
=======================================
  Files          34       34           
  Lines        3539     3539           
=======================================
  Hits         3401     3401           
  Misses        138      138           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[CUDA healpix fft] Bug when using a method="jax_cuda" with healpix sampling under shard_map

1 participant