Fix pyright failures from numpy 2.

patrick-kidger · Dec 19, 2024 · 3553ae1 · 3553ae1
1 parent 613dd0e
commit 3553ae1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/diffrax/_misc.py b/diffrax/_misc.py
@@ -148,7 +148,7 @@ def static_select(pred: BoolScalarLike, a: ArrayLike, b: ArrayLike) -> ArrayLike
     # This in turn allows us to perform some trace-time optimisations that XLA isn't
     # smart enough to do on its own.
     if isinstance(pred, (np.ndarray, np.generic)) and pred.shape == ():
-        pred = pred.item()
+        pred = cast(BoolScalarLike, pred.item())
     if pred is True:
         return a
     elif pred is False:

diff --git a/diffrax/_progress_meter.py b/diffrax/_progress_meter.py
@@ -123,8 +123,9 @@ def _step_bar(bar: list[float], progress: FloatScalarLike) -> None:
         if eqx.is_array(progress):
             # May not be an array when called with `JAX_DISABLE_JIT=1`
             progress = cast(Union[Array, np.ndarray], progress)
-            progress = progress.item()
-        progress = cast(float, progress)
+            progress = cast(float, progress.item())
+        else:
+            progress = cast(float, progress)
         bar[0] = progress
         print(f"{100 * progress:.2f}%")