jax

d6456448 - [0.6.0-UT] Skip test_dot for ROCm <6.5, and shapes exceed shared memory limits (#448)