Fix for BFGS + diffeqsolve

See patrick-kidger/optimistix#67
patrick-kidger · Jul 7, 2024 · 23b3042 · 23b3042
1 parent 70d664c
commit 23b3042
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/equinox/internal/_loop/checkpointed.py b/equinox/internal/_loop/checkpointed.py
@@ -767,7 +767,11 @@ def _checkpointed_while_loop_bwd(
     )
     del cond_fun, perturbed
     num_steps, init_residual_steps, init_residuals, filled_buffers = remainders
-    num_steps, init_residual_steps = nonbatchable((num_steps, init_residual_steps))
+    # `allow_constant_across_batch` to handle
+    # https://github.com/patrick-kidger/optimistix/issues/67
+    num_steps, init_residual_steps = nonbatchable(
+        (num_steps, init_residual_steps), allow_constant_across_batch=True
+    )
 
     def _cond_fun(carry):
         _, step_grad_val, *_ = carry

diff --git a/equinox/internal/_nontraceable.py b/equinox/internal/_nontraceable.py
@@ -131,7 +131,7 @@ def _cannot_batch(x, b, *, msg, allow_constant_across_batch):
         return x, b
     else:
         if allow_constant_across_batch:
-            x = error_if(x, jnp.min(x) != jnp.max(x), msg)
+            x = error_if(x, jnp.min(x, axis=b) != jnp.max(x, axis=b), msg)
             return jnp.take(x, 0, axis=b), batching.not_mapped
         else:
             raise ValueError(msg)