jax

3f8df9d2 - [ROCm] update to test if there are GPU events when doing profiling on ROCm via matmul