jax

417da5d0 - Add env flag to fail when testing collective matmul without nvshmem support.