/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.

Primary LanguageDockerfileMIT LicenseMIT

No issues in this repository yet.