Fix for Optimistix 0.1.0

Author: patrick-kidger

diffrax/_root_finder/_verychord.py

@@ -162,7 +162,7 @@ def terminate(
         converged = _converged(factor, self.kappa)
         terminate = at_least_two & (small | diverged | converged)
         terminate_result = optx.RESULTS.where(
-            jnp.invert(small) & (diverged | jnp.invert(converged)),
+            at_least_two & jnp.invert(small) & (diverged | jnp.invert(converged)),
             optx.RESULTS.nonlinear_divergence,
             optx.RESULTS.successful,
         )

test/test_very_chord.py

@@ -24,6 +24,9 @@ def _fn2(x, args):
 @jax.jit
 def _fn3(x, args):
     mlp = eqx.nn.MLP(4, 4, 256, 2, key=jr.PRNGKey(678))
+    dynamic, static = eqx.partition(mlp, eqx.is_array)
+    dynamic = jtu.tree_map(lambda x: x * 0.1, dynamic)
+    mlp = eqx.combine(dynamic, static)
     return mlp(x) - x