Work around JAX issue in 0.4.29

patrick-kidger · Jun 14, 2024 · aca5a16 · aca5a16
1 parent 5e57351
commit aca5a16
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
@@ -144,7 +144,8 @@ def _check(term_cls, term, term_contr_kwargs, yi):
                     raise ValueError
 
                 contr = ft.partial(term.contr, **term_contr_kwargs)
-                control_type = jax.eval_shape(contr, 0.0, 0.0)
+                # Work around https://github.com/google/jax/issues/21825
+                control_type = eqx.filter_eval_shape(contr, 0.0, 0.0)
                 control_type_compatible = eqx.filter_eval_shape(
                     better_isinstance, control_type, control_type_expected
                 )

diff --git a/diffrax/_term.py b/diffrax/_term.py
@@ -493,7 +493,7 @@ def is_vf_expensive(
         ],
         args: Args,
     ) -> bool:
-        control_struct = jax.eval_shape(self.contr, t0, t1)
+        control_struct = eqx.filter_eval_shape(self.contr, t0, t1)
         if sum(c.size for c in jtu.tree_leaves(control_struct)) in (0, 1):
             return False
         else: