JAX: Difference between revisions
imported>Nix m add Software/Applications subcategory |
imported>Samuela No edit summary |
||
| Line 1: | Line 1: | ||
[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib}</code> packages. | [https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib, jaxlibWithCuda}</code> packages. | ||
NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!). | NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!). | ||
| Line 33: | Line 33: | ||
python3 | python3 | ||
python3Packages.jax | python3Packages.jax | ||
python3Packages.jaxlibWithCuda | |||
]; | ]; | ||
} | } | ||
| Line 42: | Line 42: | ||
It should print either "cpu", "gpu", or "tpu". | It should print either "cpu", "gpu", or "tpu". | ||
Note that hydra may not cache `jaxlibWithCuda` builds on cache.nixos.org since CUDA is "unfree." @samuela publishes builds on a public cachix [https://app.cachix.org/cache/ploop#pull ploop] cache. These are periodically built and pushed from [https://github.com/samuela/nixpkgs-upkeep/ nixpkgs-upkeep]. | |||
== FAQ == | == FAQ == | ||