jax
bb524b06 - Fix KeyError for bytes_reservable_limit on ROCm

Commit

31 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

disable-inline-asm-test-jaxlib-v0.8.0

#593 - Fix KeyError for bytes_reservable_limit on ROCm

#670 - Test pytest abort plugin

#701 - Skip test_tridiagonal_solve_grad test 0.8.0

Author

magaonka-amd

Committer

Ruturaj4

Parents

894ba31f

jax bb524b06 - Fix KeyError for bytes_reservable_limit on ROCm

jax
bb524b06 - Fix KeyError for bytes_reservable_limit on ROCm