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.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Jax TPU VM Docker

Run Docker Container

sudo docker run --net=host --privileged -it --entrypoint=/bin/bash orenleung/jax_tpu:latest

Build Docker Image

sudo docker build -t orenleung/jax_tpu .

About

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.


Languages

Language:Dockerfile 100.0%