JAX: Difference between revisions

imported>Samuela
No edit summary
imported>Samuela
No edit summary
Line 2: Line 2:


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!).  
NOTE: JAX has not yet landed in nixpkgs master (see https://github.com/NixOS/nixpkgs/pull/134894). However, it is available via a fork: https://github.com/samuela/nixpkgs/tree/scratch.


== Example shell.nix, CPU only ==
== Example shell.nix, CPU only ==
<syntaxHighlight lang=nix>
<syntaxHighlight lang=nix>
let
let
  # Fork with jax/jaxlib. See https://github.com/NixOS/nixpkgs/pull/134894.
   pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/eac6215607e4ccceb9631b01ee8f8903a6e82e02.tar.gz")) {};
   pkgs = import (fetchTarball("https://github.com/samuela/nixpkgs/archive/ce58a33c2d5ff9f6cf295c6dc3fb439a467eec79.tar.gz")) {};
in pkgs.mkShell {
in pkgs.mkShell {
   buildInputs = with pkgs; [
   buildInputs = with pkgs; [
Line 31: Line 28:
let
let
   # Fork with jax/jaxlib. See https://github.com/NixOS/nixpkgs/pull/134894.
   # Fork with jax/jaxlib. See https://github.com/NixOS/nixpkgs/pull/134894.
   pkgs = import (fetchTarball("https://github.com/samuela/nixpkgs/archive/ce58a33c2d5ff9f6cf295c6dc3fb439a467eec79.tar.gz")) {};
   pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/eac6215607e4ccceb9631b01ee8f8903a6e82e02.tar.gz")) {};
in pkgs.mkShell {
in pkgs.mkShell {
   buildInputs = with pkgs; [
   buildInputs = with pkgs; [