python3Packages.jax: Include only jax.* in the built wheel

Signed-off-by: Anders Kaseorg <andersk@mit.edu>

pkgs/development/python-modules/jax/default.nix

@@ -6,6 +6,7 @@
   lapack,
   buildPythonPackage,
   fetchFromGitHub,
+  fetchpatch,
   cudaSupport ? config.cudaSupport,
 
   # build-system
@@ -52,6 +53,15 @@ buildPythonPackage (finalAttrs: {
     hash = "sha256-/RCihrjONN/+QwyQRNEmlIa7JsCLzz+SkBe5sd+ThgU=";
   };
 
+  patches = [
+    # setup.py: Include only jax.* in the built wheel
+    # https://github.com/jax-ml/jax/pull/37182
+    (fetchpatch {
+      url = "https://github.com/jax-ml/jax/commit/cb5a91780f84f124090d8f94e99c8771e87590f6.patch";
+      hash = "sha256-p6X9IFe4YUb2MQp7YjJHme1dueZ1Y37IKnANGruW1cM=";
+    })
+  ];
+
   build-system = [ setuptools ];
 
   # The version is automatically set to ".dev" if this variable is not set.