From e87d72216a12e3882deba5da05a18adce6f776ba Mon Sep 17 00:00:00 2001
From: Anders Kaseorg <andersk@mit.edu>
Date: Sat, 25 Apr 2026 14:09:55 -0700
Subject: [PATCH] python3Packages.jax: Include only jax.* in the built wheel

Signed-off-by: Anders Kaseorg <andersk@mit.edu>
---
 pkgs/development/python-modules/jax/default.nix | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
index d2a0a7109cf2..a8a84ae2747b 100644
--- a/pkgs/development/python-modules/jax/default.nix
+++ b/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.