fixup! Fixed a major source of bugs: ControlTerms no longer broadcast.

patrick-kidger · Jan 13, 2025 · 41b3cf1 · 41b3cf1
1 parent a828e5b
commit 41b3cf1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/diffrax/_term.py b/diffrax/_term.py
@@ -404,6 +404,9 @@ def vector_field(t, y, args):
         ```
     """  # noqa: E501
 
+    vector_field: Callable[[RealScalarLike, Y, Args], _VF]
+    control: AbstractPath[_Control]
+
     def __init__(
         self,
         vector_field: Callable[[RealScalarLike, Y, Args], _VF],

diff --git a/test/test_term.py b/test/test_term.py
@@ -85,7 +85,8 @@ def derivative(self, t, left=True):
             return jr.normal(derivkey, (3,))
 
     control = Control()
-    term = diffrax.WeaklyDiagonalControlTerm(vector_field, control)
+    with pytest.warns(match="`WeaklyDiagonalControlTerm` is now deprecated"):
+        term = diffrax.WeaklyDiagonalControlTerm(vector_field, control)
     args = getkey()
     dx = term.contr(0, 1)
     y = jnp.array([1.0, 2.0, 3.0])