Cuda and merge upstream jax flake