Skip to content

Skip test_tridiagonal_solve_grad test 0.8.0#701

Open
AratiGanesh wants to merge 1 commit intorocm-jaxlib-v0.8.0from
skip-tridiagonal-solve-grad-v0.8.0
Open

Skip test_tridiagonal_solve_grad test 0.8.0#701
AratiGanesh wants to merge 1 commit intorocm-jaxlib-v0.8.0from
skip-tridiagonal-solve-grad-v0.8.0

Conversation

@AratiGanesh
Copy link

@AratiGanesh AratiGanesh commented Feb 9, 2026

Motivation

The test_tridiagonal_solve_grad test fails on ROCm devices due to numerical precision issues in the underlying rocSparse library.

Technical Details

  • Added conditional skip logic using jtu.is_device_rocm(), shape == (3, 4), and dtype == np.float32 to skip only test_tridiagonal_solve_grad0 on ROCm in tests/linalg_test.py, preserving other passing test variants.

Test Result

image

Upstream PR - jax-ml#34959

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

@AratiGanesh AratiGanesh changed the title Skip test_tridiagonal_solve_grad on ROCm due to rocSparse issue Skip test_tridiagonal_solve_grad test 0.8.0 Feb 9, 2026
@AratiGanesh AratiGanesh force-pushed the skip-tridiagonal-solve-grad-v0.8.0 branch from 50088a5 to fe8d3e9 Compare February 10, 2026 15:45
@AratiGanesh AratiGanesh added cherry-pick-candidate Mark a PR to be cherry-picked into the next ROCm JAX. Remove IIF the latest upstream contain the PR. open-upstream Tag when you want a copy of this PR to be opened on upstream labels Feb 10, 2026
@AratiGanesh AratiGanesh requested a review from a team February 10, 2026 23:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cherry-pick-candidate Mark a PR to be cherry-picked into the next ROCm JAX. Remove IIF the latest upstream contain the PR. open-upstream Tag when you want a copy of this PR to be opened on upstream

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant