Fix newly exposed mypy issues

benchmark/plotters.py

@@ -20,6 +20,7 @@
 
 import pandas as pd
 from matplotlib.axes import Axes
+from matplotlib.dates import date2num
 
 from benchmark.grouping import GroupingKey, GroupingSpec, group
 from benchmark.metadata import BenchmarkMetadata, parse_timestamp
@@ -118,7 +119,7 @@ def time_line(
         timestamp_by = GroupingSpec([GroupingKey.TIMESTAMP], minimise=False)
         timestamp_groups = group(df, [], metadata, timestamp_by)
         for (timestamp,), timestamp_df, _ in timestamp_groups:
-            parsed_timestamp = parse_timestamp(timestamp)
+            parsed_timestamp = date2num(parse_timestamp(timestamp))
             line_xs.append(parsed_timestamp)
             y_mean = timestamp_df.value.mean()
             y_std = timestamp_df.value.std()

gpflow/base.py

@@ -27,19 +27,19 @@
 if TYPE_CHECKING:  # pragma: no cover
     from IPython.lib import pretty
 
-DType = Union[np.dtype, tf.DType]
-
-
 if TYPE_CHECKING and (not NP_TYPE_CHECKING):  # pragma: no cover
     AnyNDArray = Any
+    DType = Any
 else:
     if GENERIC_NP_ARRAYS:
         # It would be nice to use something more interesting than `Any` here, but it looks like
         # the infrastructure in the rest of the ecosystem isn't really set up for this
         # yet. Maybe when we get Python 3.11?
         AnyNDArray = np.ndarray[Any, Any]  # type: ignore[misc]
+        DType = Union[np.dtype[Any], tf.DType]  # type: ignore[misc]
     else:
         AnyNDArray = Union[np.ndarray]  # type: ignore[misc]
+        DType = Union[np.dtype, tf.DType]  # type: ignore[misc]
 
 VariableData = Union[List[Any], Tuple[Any], AnyNDArray, int, float]  # deprecated
 Transform = Union[tfp.bijectors.Bijector]

gpflow/monitor/tensorboard.py

@@ -20,7 +20,7 @@
 import numpy as np
 import tensorflow as tf
 
-from ..base import Parameter
+from ..base import AnyNDArray, Parameter
 from ..models import BayesianModel
 from ..utilities import parameter_dict
 from .base import MonitorTask
@@ -168,7 +168,11 @@ class ImageToTensorBoard(ToTensorBoard):
     def __init__(
         self,
         log_dir: str,
-        plotting_function: Callable[["matplotlib.figure.Figure", "matplotlib.figure.Axes"], None],
+        plotting_function: Union[
+            # if len(subplots_kw) > 1 then the function is passed an array of axes
+            Callable[["matplotlib.figure.Figure", "matplotlib.axes.Axes"], None],
+            Callable[["matplotlib.figure.Figure", AnyNDArray], None],
+        ],
         name: Optional[str] = None,
         *,
         fig_kw: Optional[Dict[str, Any]] = None,

gpflow/optimizers/natgrad.py

@@ -16,13 +16,12 @@
 import functools
 from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
 
-import numpy as np
 import tensorflow as tf
 from check_shapes import check_shapes
 
-from ..base import Parameter, _to_constrained
+from ..base import AnyNDArray, Parameter, _to_constrained
 
-Scalar = Union[float, tf.Tensor, np.ndarray]
+Scalar = Union[float, tf.Tensor, AnyNDArray]
 LossClosure = Callable[[], tf.Tensor]
 NatGradParameters = Union[Tuple[Parameter, Parameter], Tuple[Parameter, Parameter, "XiTransform"]]
 

tests/gpflow/test_monitor.py

@@ -11,6 +11,7 @@
 from matplotlib.figure import Figure
 
 import gpflow
+from gpflow.base import AnyNDArray
 from gpflow.models import GPR, GPModel
 from gpflow.monitor import (
     ExecuteCallback,
@@ -105,7 +106,7 @@ def test_ImageToTensorBoard(tmp_path: Path) -> None:
     """Smoke test `ImageToTensorBoard` in Eager and Compiled mode."""
     tmp_path_str = str(tmp_path)
 
-    def plotting_cb(fig: Figure, axes: Axes) -> None:
+    def plotting_cb(fig: Figure, axes: AnyNDArray) -> None:
         axes[0, 0].plot(np.random.randn(2), np.random.randn(2))
         axes[1, 0].plot(np.random.randn(2), np.random.randn(2))
         axes[0, 1].plot(np.random.randn(2), np.random.randn(2))