eden-chan / jax_tpu_docker Goto Github PK
View Code? Open in Web Editor NEWThis project forked from davidnet/jax_tpu_docker
Provides a Docker image for running JAX on TPU VMs. This repository contains the necessary files to build the Docker image and run JAX on TPUs with ease, making it easy for researchers to take advantage of the power of TPUs for their machine learning projects.
License: MIT License