Merge pull request #134894 from samuela/samuela/jax

4 files changed, 139 insertions, 0 deletions

diff --git a/maintainers/maintainer-list.nix b/maintainers/maintainer-list.nix
index 04033780a62..b945ddc83c5 100644
--- a/maintainers/maintainer-list.nix
+++ b/maintainers/maintainer-list.nix
@@ -9594,6 +9594,12 @@
     githubId = 115821;
     name = "Sam Rose";
   };
+  samuela = {
+    email = "skainsworth@gmail.com";
+    github = "samuela";
+    githubId = 226872;
+    name = "Samuel Ainsworth";
+  };
   samueldr = {
     email = "samuel@dionne-riel.com";
     github = "samueldr";
diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
new file mode 100644
index 00000000000..f0b4d6e2632
--- /dev/null
+++ b/pkgs/development/python-modules/jax/default.nix
@@ -0,0 +1,39 @@
+{ buildPythonPackage, fetchFromGitHub, lib
+# propagatedBuildInputs
+, absl-py, numpy, opt-einsum
+# checkInputs
+, jaxlib, pytestCheckHook
+}:
+
+buildPythonPackage rec {
+  pname = "jax";
+  version = "0.2.19";
+
+  # Fetching from pypi doesn't allow us to run the test suite. See https://discourse.nixos.org/t/pythonremovetestsdir-hook-being-run-before-checkphase/14612/3.
+  src = fetchFromGitHub {
+    owner = "google";
+    repo = pname;
+    rev = "jax-v${version}";
+    sha256 = "sha256-pVn62G7pydR7ybkf7gSbu0FlEq2c0US6H2GTBAljup4=";
+  };
+
+  # jaxlib is _not_ included in propagatedBuildInputs because there are
+  # different versions of jaxlib depending on the desired target hardware. The
+  # JAX project ships separate wheels for CPU, GPU, and TPU. Currently only the
+  # CPU wheel is packaged.
+  propagatedBuildInputs = [ absl-py numpy opt-einsum ];
+
+  checkInputs = [ jaxlib pytestCheckHook ];
+  # NOTE: Don't run the tests in the expiremental directory as they require flax
+  # which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2.
+  # Not a big deal, this is how the JAX docs suggest running the test suite
+  # anyhow.
+  pytestFlagsArray = [ "-W ignore::DeprecationWarning" "tests/" ];
+
+  meta = with lib; {
+    description = "Differentiate, compile, and transform Numpy code";
+    homepage    = "https://github.com/google/jax";
+    license     = licenses.asl20;
+    maintainers = with maintainers; [ samuela ];
+  };
+}
diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix
new file mode 100644
index 00000000000..7761c037299
--- /dev/null
+++ b/pkgs/development/python-modules/jaxlib/default.nix
@@ -0,0 +1,90 @@
+# For the moment we only support the CPU and GPU backends of jaxlib. The TPU
+# backend will require some additional work. Those wheels are located here:
+# https://storage.googleapis.com/jax-releases/libtpu_releases.html.
+
+# For future reference, the easiest way to test the GPU backend is to run
+#   NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib.override { cudaSupport = true; }"
+#   export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1
+#   python -c "from jax.lib import xla_bridge; assert xla_bridge.get_backend().platform == 'gpu'"
+#   python -c "from jax import random; random.PRNGKey(0)"
+#   python -c "from jax import random; x = random.normal(random.PRNGKey(0), (100, 100)); x @ x"
+# There's no convenient way to test the GPU backend in the derivation since the
+# nix build environment blocks access to the GPU. See also:
+#   * https://github.com/google/jax/issues/971#issuecomment-508216439
+#   * https://github.com/google/jax/issues/5723#issuecomment-913038780
+
+{ addOpenGLRunpath, autoPatchelfHook, buildPythonPackage, config, fetchPypi
+, fetchurl, isPy39, lib, stdenv
+# propagatedBuildInputs
+, absl-py, flatbuffers, scipy, cudatoolkit_11
+# Options:
+, cudaSupport ? config.cudaSupport or false
+}:
+
+assert cudaSupport -> lib.versionAtLeast cudatoolkit_11.version "11.1";
+
+let
+  device = if cudaSupport then "gpu" else "cpu";
+in
+buildPythonPackage rec {
+  pname = "jaxlib";
+  version = "0.1.71";
+  format = "wheel";
+
+  # At the time of writing (8/19/21), there are releases for 3.7-3.9. Supporting
+  # all of them is a pain, so we focus on 3.9, the current nixpkgs python3
+  # version.
+  disabled = !isPy39;
+
+  src = {
+    cpu = fetchurl {
+      url = "https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp39-none-manylinux2010_x86_64.whl";
+      sha256 = "sha256:0rqhs6qabydizlv5d3rb20dbv6612rr7dqfniy9r6h4kazdinsn6";
+    };
+    gpu = fetchurl {
+      url = "https://storage.googleapis.com/jax-releases/cuda111/jaxlib-${version}+cuda111-cp39-none-manylinux2010_x86_64.whl";
+      sha256 = "sha256:065kyzjsk9m84d138p99iymdiiicm1qz8a3iwxz8rspl43rwrw89";
+    };
+  }.${device};
+
+  # Prebuilt wheels are dynamically linked against things that nix can't find.
+  # Run `autoPatchelfHook` to automagically fix them.
+  nativeBuildInputs = [ autoPatchelfHook ] ++ lib.optional cudaSupport addOpenGLRunpath;
+  # Dynamic link dependencies
+  buildInputs = [ stdenv.cc.cc ];
+
+  # jaxlib contains shared libraries that open other shared libraries via dlopen
+  # and these implicit dependencies are not recognized by ldd or
+  # autoPatchelfHook. That means we need to sneak them into rpath. This step
+  # must be done after autoPatchelfHook and the automatic stripping of
+  # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
+  # patchPhase. Dependencies:
+  #   * libcudart.so.11.0 -> cudatoolkit_11.lib
+  #   * libcublas.so.11   -> cudatoolkit_11
+  #   * libcuda.so.1      -> opengl driver in /run/opengl-driver/lib
+  preInstallCheck = lib.optional cudaSupport ''
+    shopt -s globstar
+
+    addOpenGLRunpath $out/**/*.so
+
+    for file in $out/**/*.so; do
+      rpath=$(patchelf --print-rpath $file)
+      # For some reason `makeLibraryPath` on `cudatoolkit_11` maps to
+      # <cudatoolkit_11.lib>/lib which is different from <cudatoolkit_11>/lib.
+      patchelf --set-rpath "$rpath:${cudatoolkit_11}/lib:${lib.makeLibraryPath [ cudatoolkit_11.lib ]}" $file
+    done
+  '';
+
+  # pip dependencies and optionally cudatoolkit.
+  propagatedBuildInputs = [ absl-py flatbuffers scipy ] ++ lib.optional cudaSupport cudatoolkit_11;
+
+  pythonImportsCheck = [ "jaxlib" ];
+
+  meta = with lib; {
+    description = "XLA library for JAX";
+    homepage    = "https://github.com/google/jax";
+    license     = licenses.asl20;
+    maintainers = with maintainers; [ samuela ];
+    platforms = [ "x86_64-linux" ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index ecb2bf3e0da..e99f277c7e0 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -3732,6 +3732,10 @@ in {
 
   javaproperties = callPackage ../development/python-modules/javaproperties { };
 
+  jax = callPackage ../development/python-modules/jax { };
+
+  jaxlib = callPackage ../development/python-modules/jaxlib { };
+
   JayDeBeApi = callPackage ../development/python-modules/JayDeBeApi { };
 
   jc = callPackage ../development/python-modules/jc { };