minor fixes

diffrax/solver/additive_srk.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Optional
 
 import equinox as eqx
 import jax
@@ -13,7 +13,7 @@
 from ..custom_types import Array, Bool, DenseInfo, LevyVal, PyTree, Scalar
 from ..local_interpolation import LocalLinearInterpolation
 from ..solution import RESULTS
-from ..term import _ControlTerm, AbstractTerm, MultiTerm, ODETerm
+from ..term import AbstractTerm, MultiTerm, ODETerm
 from .base import AbstractStratonovichSolver
 
 
@@ -53,8 +53,6 @@ def __post_init__(self):
             assert np.allclose(sum(a_i), c_i)
         assert np.allclose(sum(self.b_sol), 1.0)
 
-        # TODO: add checks for whether the method is FSAL
-
 
 StochasticButcherTableau.__init__.__doc__ = """**Arguments:**
 
@@ -111,13 +109,13 @@ class AbstractAdditiveSRK(AbstractStratonovichSolver):
     in the `StochasticButcherTableau`.
     """
 
-    term_structure = MultiTerm[Tuple[ODETerm, _ControlTerm]]
+    term_structure = MultiTerm[tuple[ODETerm, AbstractTerm]]
     interpolation_cls = LocalLinearInterpolation
     tableau: StochasticButcherTableau
 
     def init(
         self,
-        terms: MultiTerm[Tuple[ODETerm, _ControlTerm]],
+        terms: term_structure,
         t0: Scalar,
         t1: Scalar,
         y0: PyTree,
@@ -132,7 +130,7 @@ def init(
             if not path.levy_area == "space-time":
                 raise ValueError(
                     "The Brownian path controlling the diffusion "
-                    "should be initialised with `compute_stla=True`"
+                    "should be initialised with `levy_area='space-time'`"
                 )
 
         # check that the vector field of the diffusion term is constant
@@ -157,14 +155,14 @@ def _embed_a_lower(self, dtype):
 
     def step(
         self,
-        terms: MultiTerm[Tuple[ODETerm, _ControlTerm]],
+        terms: term_structure,
         t0: Scalar,
         t1: Scalar,
         y0: PyTree,
         args: PyTree,
         solver_state: _SolverState,
         made_jump: Bool,
-    ) -> Tuple[PyTree, _ErrorEstimate, DenseInfo, _SolverState, RESULTS]:
+    ) -> tuple[PyTree, _ErrorEstimate, DenseInfo, _SolverState, RESULTS]:
         del solver_state, made_jump
 
         h = t1 - t0