The toolkit in the packages by default is old though. I was setting up JAX on a Ubuntu 20.04 WSL 2 install and it was complaining about ptxas being old and giving bad results.