Skip to content

Add ROCm gesdd SVD backend, export back-compat data, and relaxed GPU …#734

Open
cj401-amd wants to merge 1 commit intorocm-jaxlib-v0.9.0from
ci_cj_gesdd_rocm-jaxlib-v0.9.0
Open

Add ROCm gesdd SVD backend, export back-compat data, and relaxed GPU …#734
cj401-amd wants to merge 1 commit intorocm-jaxlib-v0.9.0from
ci_cj_gesdd_rocm-jaxlib-v0.9.0

Conversation

@cj401-amd
Copy link

@cj401-amd cj401-amd added the enhancement New feature or request label Mar 6, 2026
@tsrw2048
Copy link

tsrw2048 commented Mar 6, 2026

@cj401-amd Thank you for adding the GESDD call! The additional functionality is much appreciated. I had a few requests regarding the PR itself.

  • Would you mind breaking up this PR into several? Right now the description just covers GESDD but the code changes cover GESDD, export tests for other FFI calls, and changes to other lowering rules like GEQRF and LU. In each PR, we need to know the reason for the change and associated tests for it. This makes it easy to track changes from one release to the next as well as debug issues in the future. For example, if this PR could contain just the changes relevant for the GESDD implementation (similar to what you did upstream) that would make it much easier to review.
  • Some of these changes (such as modifications to the lowering rules) are from upstream commits. When possible, can you cherry pick the commits from upstream and include the attribution to the original commit/author. This helps us track where and why upstream made these changes in the event that something breaks downstream. Also, please don't include the upstream changes if they are not needed for the PR itself. We will get these changes naturally when we move from release to release.
  • I think you can leave out everything related to the export testdata and back compat tests. We have intentionally left out the export testdata files and the associated unit tests in export_back_compat_test_data.py. The test data files are baked into the JAX wheel and imported into the unit test file. Since we are not building the JAX wheel in the downstream CI, this will cause a bunch of failures based on import errors. To solve this, we have pushed all the test data files and the back compat unit tests directly upstream (and they have already been merged).

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

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants