maint: disable new jax jit implementation

lsqfitgp/_patch_jax.py

@@ -20,6 +20,11 @@
 import pathlib
 import traceback
 
+# disable this jax 0.4.4 feature because it makes empbayes_fit fail under jit
+# due to leaked tracers
+import os
+os.environ['JAX_JIT_PJIT_API_MERGE'] = '0'
+
 from jax.config import config
 config.update("jax_enable_x64", True)
 
@@ -120,9 +125,12 @@ def __exit__(self, exc_type, exc_value, tb):
         if exit or exc_type in (jax.errors.ConcretizationTypeError, jax.errors.TracerArrayConversionError):
             return True
 
-        if exc_type is IndexError and traceback.extract_tb(tb)[-1].name == 'arg_info_pytree': # pragma: no cover
+        weird_cond = exc_type is IndexError and (
+            traceback.extract_tb(tb)[-1].name in ('arg_info_pytree', '_origin_msg'),
+        )
+        if weird_cond: # pragma: no cover
             # TODO this ignores a jax internal bug I don't understand, appears
-            # in examples/pdf4.py, and not on CI
+            # in examples/pdf4.py
             return True
 
 # TODO make stop_hessian work in reverse mode