elixir-nx / xla

Pre-compiled XLA extension
Apache License 2.0
83 stars 21 forks source link

Streamline docker builds #50

Closed jonatanklosko closed 10 months ago

jonatanklosko commented 10 months ago

This simplifies the flow to build CUDA binaries on a remote server.

I also realised that we can install a specific cuDNN version in the CUDA Docker container (since the container usually has the latest). This way we don't have to point to old image SHAs, which is especially tricky because there is no way to find old SHAs for a given tag.

The CUDA binaries from v0.5.0 are build with cuDNN 8.9, which is more recent than we want. I'm almost finished building with older cuDNN and will update the binaries in place :)