Request for 🐳 Dockerfile

[SauravMaheshkar]

Hey folks, I just came across the repository and wanted to highlight a use-case. Getting JAX up and running with the proper CUDA and CUDNN versions can be a hassle and IMO Docker is the quickest workaround for easy onboarding and setup. It'd be much easier if there was a Dockerfile provided or even an Image on Docker-hub or the Github Container Registry (ghcr.io). That would make reproducibility and benchmarking much easier. At the moment, JAX is a bit spotty with their Docker configurations, but the tensorflow-gpu image works perfectly as a base and already contains many of the needed packages. Having a Dockerfile could also help in creating and writing slightly compute heavy examples with proper accelerator (GPU/TPU) setup.