jax
cef02f3a - Fix KeyError for bytes_reservable_limit on ROCm

Commit

26 days ago

Fix KeyError for bytes_reservable_limit on ROCm ROCm's memory_stats() returns 'bytes_limit' instead of 'bytes_reservable_limit'. Add platform-specific handling to test_lax_full_like_efficient to support both CUDA and ROCm.

References

#609 - Cherry 8.2

#632 - Add script for running bazel tests on ROCm

#638 - Support wheels as build deps

#704 - Skip test_tridiagonal_solve_grad test 0.8.2

#706 - Backporting `device test` unit tests for ROCm from JAX v0.9.0 to v0.8.2.

#707 - Skip test_batch_axis_sharding_jvp13 test on 0.8.2

#710 - fix poor performance of slogdet when input dimension n=1 or 2

Author

magaonka-amd

Committer

Ruturaj4

Parents

e3ffcb08

jax cef02f3a - Fix KeyError for bytes_reservable_limit on ROCm

jax
cef02f3a - Fix KeyError for bytes_reservable_limit on ROCm