jax

821193b2 - Explicitly build specific CUDA capabilities. (#2722)

Explicitly build specific CUDA capabilities. (#2722) We choose the same set as TensorFlow (minus 3.7, which TF is apparently considering dropping anyway). This avoids a slow PTX -> SASS compilation on first time startup.