From 1f8686373abe21bf6b1ce972f0ea55405b449329 Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Sun, 22 Aug 2021 20:37:42 +0000
Subject: [PATCH] python3Packages.jaxlib: init at 0.1.70

---
 .../python-modules/jaxlib/default.nix         | 45 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 47 insertions(+)
 create mode 100644 pkgs/development/python-modules/jaxlib/default.nix

diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix
new file mode 100644
index 000000000000..240c5a7d6d0e
--- /dev/null
+++ b/pkgs/development/python-modules/jaxlib/default.nix
@@ -0,0 +1,45 @@
+# For the moment we only support the CPU backend of jaxlib. GPU and TPU backends
+# require some additional work. Their wheels are not located on PyPI.
+#  * CPU/GPU: https://storage.googleapis.com/jax-releases/jax_releases.html
+#  * TPU: https://storage.googleapis.com/jax-releases/libtpu_releases.html
+
+{ autoPatchelfHook, buildPythonPackage, fetchPypi, isPy39, lib, stdenv
+# propagatedBuildInputs
+, absl-py, flatbuffers, scipy
+}:
+
+buildPythonPackage rec {
+  pname = "jaxlib";
+  version = "0.1.70";
+  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 = fetchPypi {
+    inherit pname version format;
+    dist = "cp39";
+    python = "cp39";
+    platform = "manylinux2010_x86_64";
+    sha256 = "sha256-mytMTqoavpuRawj52MU5/iFj27SGlm8DaoQ5vd/3bss=";
+  };
+
+  # Prebuilt wheels are dynamically linked against things that nix can't find.
+  # Run `autoPatchelfHook` to automagically fix them.
+  nativeBuildInputs = [ autoPatchelfHook ];
+  # Dynamic link dependencies
+  buildInputs = [ stdenv.cc.cc ];
+
+  # pip dependencies
+  propagatedBuildInputs = [ absl-py flatbuffers scipy ];
+
+  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 428913333280..fa17a8a7e7ee 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -3716,6 +3716,8 @@ in {
 
   javaproperties = callPackage ../development/python-modules/javaproperties { };
 
+  jaxlib = callPackage ../development/python-modules/jaxlib { };
+
   JayDeBeApi = callPackage ../development/python-modules/JayDeBeApi { };
 
   jc = callPackage ../development/python-modules/jc { };