int

Author: lockwo

diffrax/_integrate.py

@@ -187,10 +187,10 @@ def _check(term_cls, term, term_contr_kwargs, yi):
                 contr = ft.partial(term.contr, **term_contr_kwargs)
                 # Work around https://github.com/google/jax/issues/21825
                 try:
-                    control_type = eqx.filter_eval_shape(contr, t, t)
+                    control_type, path_type = eqx.filter_eval_shape(contr, t, t)
                 except Exception as e:
                     raise ValueError(f"Error while tracing {term}.contr: " + str(e))
-                control_type_compatible = eqx.filter_eval_shape(
+                control_type_compatible, path_type_expected = eqx.filter_eval_shape(
                     better_isinstance, control_type, control_type_expected
                 )
                 if not control_type_compatible: