_integrate.py

patrick-kidger · Jan 14, 2025 · d96572e · d96572e
1 parent 134a40a
commit d96572e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
@@ -649,7 +649,8 @@ def body_fun(state):
         event_mask = final_state.event_mask
         flat_mask = jtu.tree_leaves(event_mask)
         assert all(jnp.shape(x) == () for x in flat_mask)
-        event_happened = jnp.any(jnp.stack(flat_mask))
+        float_mask = jnp.array(flat_mask).astype(jnp.float32)
+        event_happened = jnp.max(float_mask) > 0.0
 
         def _root_find():
             _interpolator = solver.interpolation_cls(