_control term isn't recognized for some reason

Author: lockwo

diffrax/_solver/kl.py

@@ -1,5 +1,5 @@
 import operator
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import equinox as eqx
 import jax
@@ -20,8 +20,8 @@
 from .._heuristics import is_sde
 from .._solution import RESULTS
 from .._term import (
-    _ControlTerm,
     AbstractTerm,
+    ControlTerm,
     MultiTerm,
     ODETerm,
     WeaklyDiagonalControlTerm,
@@ -33,6 +33,9 @@
 )
 
 
+_ControlTerm = Union[ControlTerm, WeaklyDiagonalControlTerm]
+
+
 def _compute_kl_integral(
     drift_term1: ODETerm,
     drift_term2: ODETerm,