jax

0956a239 - Fix KeyError for bytes_reservable_limit on ROCm

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.