jax 0.4.25 compatibility (PjitFunction)

gboehl · Apr 8, 2024 · 9914842 · 9914842
1 parent 3c0369e
commit 9914842
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/econpizza/parser/__init__.py b/econpizza/parser/__init__.py
@@ -148,7 +148,7 @@ def _load_external_functions_file(model, context):
         module = _load_as_module(model["functions_file"])
 
         def func_or_compiled(func): return isinstance(
-            func, jaxlib.xla_extension.CompiledFunction) or isfunction(func)
+            func, jaxlib.xla_extension.PjitFunction) or isfunction(func)
         for m in getmembers(module, func_or_compiled):
             context[m[0]] = m[1]