orenleung/jax_tpu

Docker image for running JAX on TPU VM, making it easy to take advantage of the power of TPUs for ML

Latest release: over 3 years ago - 43 downloads - 16 dependencies