From cb7b514703ec672d935ef19a38238f705cec4f1c Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Sat, 16 Apr 2022 02:23:12 +0000
Subject: [PATCH] python3Packages.jax: fix build when nixpkgs-wide cudaSupport
 is enabled

---
 pkgs/development/python-modules/jax/default.nix | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
index 94be46b3aeb2..be791b26f3c0 100644
--- a/pkgs/development/python-modules/jax/default.nix
+++ b/pkgs/development/python-modules/jax/default.nix
@@ -3,6 +3,7 @@
 , blas
 , buildPythonPackage
 , fetchFromGitHub
+, fetchpatch
 , jaxlib
 , lapack
 , numpy
@@ -34,6 +35,12 @@ buildPythonPackage rec {
   patches = [
     # See https://github.com/google/jax/issues/7944
     ./cache-fix.patch
+
+    # See https://github.com/google/jax/issues/10292
+    (fetchpatch {
+      url = "https://github.com/google/jax/commit/cadc8046d56e0c1433cf48a2f106947d5f4ecbfd.patch";
+      hash = "sha256-jrpIqt4LzWAswt/Cpwtfa5d1Yn31HcXkVH3ETmaigA0=";
+    })
   ];
 
   # jaxlib is _not_ included in propagatedBuildInputs because there are