python3Packages.numpyro: fix jax 0.10.0 compatibility

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

@@ -40,6 +40,11 @@ buildPythonPackage (finalAttrs: {
     hash = "sha256-sNqllL9nBwXp0kn+HAjvIaHf7LR0UKh9q7DZ20yCr5A=";
   };
 
+  patches = [
+    # Remove usage of xla_pmap_p which was removed in jax 0.10.0
+    ./fix-jax-0.10.0-compat.patch
+  ];
+
   build-system = [ setuptools ];
 
   dependencies = [

pkgs/development/python-modules/numpyro/fix-jax-0.10.0-compat.patch

@@ -0,0 +1,20 @@
+diff --git a/numpyro/ops/provenance.py b/numpyro/ops/provenance.py
+index 1234567..abcdefg 100644
+--- a/numpyro/ops/provenance.py
++++ b/numpyro/ops/provenance.py
+@@ -4,7 +4,7 @@
+ import jax
+ from jax.api_util import debug_info, flatten_fun, shaped_abstractify
+ from jax.extend.core import Literal
+-from jax.extend.core.primitives import call_p, closed_call_p, jit_p, xla_pmap_p
++from jax.extend.core.primitives import call_p, closed_call_p, jit_p
+ import jax.extend.linear_util as lu
+ from jax.interpreters.partial_eval import trace_to_jaxpr_dynamic
+
+@@ -114,7 +114,6 @@ def track_deps_call_rule(eqn, provenance_inputs):
+
+ track_deps_rules[call_p] = track_deps_call_rule
+-track_deps_rules[xla_pmap_p] = track_deps_call_rule
+
+
+ def track_deps_closed_call_rule(eqn, provenance_inputs):