Work around JAX issue in 0.4.29

patrick-kidger · patrick-kidger · commit aca5a1643dae · 2024-06-14T15:31:00.000+02:00
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:

Original file line number	Diff line number	Diff line change
`@@ -144,7 +144,8 @@ def _check(term_cls, term, term_contr_kwargs, yi):`
`144`	`144`	`raise ValueError`
`145`	`145`
`146`	`146`	`contr = ft.partial(term.contr, **term_contr_kwargs)`
`147`		`- control_type = jax.eval_shape(contr, 0.0, 0.0)`
	`147`	`+ # Work around https://github.com/google/jax/issues/21825`
	`148`	`+ control_type = eqx.filter_eval_shape(contr, 0.0, 0.0)`
`148`	`149`	`control_type_compatible = eqx.filter_eval_shape(`
`149`	`150`	`better_isinstance, control_type, control_type_expected`
`150`	`151`	`)`