Update request: jaxlib with ROCm support

NixOS / nixpkgs

Nix Packages collection & NixOS

MIT License

16.48k stars 12.97k forks source link

Update request: jaxlib with ROCm support #302676

Open dwf opened 2 months ago

dwf commented 2 months ago

Technically JAX [and TensorFlow] support ROCm via XLA.

I've taken a run at this in https://github.com/dwf/nixpkgs/tree/jax_rocm and with a pile of hacks have gotten it to build but it fails to run a matrix multiply correctly. I used rocmPackages_5 due to #299156 blocking 6.

Previous work on TensorFlow at https://github.com/Madouura/nixpkgs/commit/344aa780809455f545eae4895bed72e3e9af0de6

cc #197885

Add a :+1: reaction to issues you find important.

dwf commented 2 months ago

Log from attempting to use the unsuccessful build (AMD_LOG_LEVEL=2 python) https://gist.github.com/dwf/f071422e3d3e09a4024a70c649b7cd51