bump version; support newer JAX versions

n-gao · n-gao · commit 34e41fac11e5 · 2025-08-23T12:41:53.000+02:00
diff --git a/folx/interpreter.py b/folx/interpreter.py
@@ -214,7 +214,7 @@ def eval_laplacian(eqn: JaxprEqn, invals):
                     outvals = eqn.primitive.bind(*subfuns, *invals, **bind_params)
             elif eqn.primitive.name == 'scan':
                 outvals = eval_scan(eqn, invals)
-            elif eqn.primitive.name == 'pjit':
+            elif eqn.primitive.name in ('jit', 'pjit'):
                 outvals = eval_pjit(eqn, invals)
             elif eqn.primitive.name == 'custom_jvp_call':
                 outvals = eval_custom_jvp(eqn, invals)
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "folx"
-version = "0.2.17"
+version = "0.2.18"
 description = "Forward Laplacian for JAX"
 authors = [
     { name = "Nicholas Gao", email = "n.gao@tum.de" },
diff --git a/uv.lock b/uv.lock
