From 65a33fec2a6ab71d165e2caacc7254c6669edec0 Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Fri, 28 Jan 2022 11:04:10 -0800
Subject: [PATCH] python3Packages.jaxlibWithCuda: init at 0.1.75 (#157055)

* python3Packages.jaxlibWithCuda: init at 0.1.75

* python3Packages.jaxlibWithoutCuda: init at 0.1.75
---
 pkgs/top-level/python-packages.nix | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 0566e7237534..99f85295f4d0 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -4127,6 +4127,14 @@ in {
 
   jaxlib = self.jaxlib-build;
 
+  jaxlibWithCuda = self.jaxlib-build.override {
+    cudaSupport = true;
+  };
+
+  jaxlibWithoutCuda = self.jaxlib-build.override {
+    cudaSupport = false;
+  };
+
   JayDeBeApi = callPackage ../development/python-modules/JayDeBeApi { };
 
   jc = callPackage ../development/python-modules/jc { };