jax
0956a239 - Fix KeyError for bytes_reservable_limit on ROCm

Commit

105 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

#664 - update with ROCm profiling tests with kernel_details

#671 - Rocm jaxlib v0.9.0

#699 - Skip testEighTinyNorm test 0.9.0

#725 - Skip test_pmap on ROCm due to IndivisibleError with new pmap SPMD tiling

#732 - update test

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

#739 - Fix: Scaled Matmul at rocm 9.0

#765 - jaxlib/gpu/rnn_kernels.cc: remove premature input_tensor_desc destruction (ROCM-21859)

Author

magaonka-amd

Committer

Ruturaj4

Parents

fc3bc807

jax 0956a239 - Fix KeyError for bytes_reservable_limit on ROCm

jax
0956a239 - Fix KeyError for bytes_reservable_limit on ROCm