gpflow 2.9.0__tar.gz → 2.9.2__tar.gz

{gpflow-2.9.0/gpflow.egg-info → gpflow-2.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpflow
-Version: 2.9.0
+Version: 2.9.2
 Summary: Gaussian process methods in TensorFlow
 Home-page: https://www.gpflow.org
 Author: James Hensman, Alex Matthews
@@ -9,6 +9,7 @@ License: Apache License 2.0
 Project-URL: Source on GitHub, https://github.com/GPflow/GPflow
 Project-URL: Documentation, https://gpflow.github.io/GPflow/
 Keywords: machine-learning gaussian-processes kernels tensorflow
+Platform: UNKNOWN
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: GPU :: NVIDIA CUDA
 Classifier: Intended Audience :: Developers
@@ -216,3 +217,4 @@ which can be cited by users.
   url = {https://arxiv.org/abs/2003.01115}
 }
 ```
+

{gpflow-2.9.0 → gpflow-2.9.2}/RELEASE.md

@@ -66,6 +66,39 @@ This release contains contributions from:
 <INSERT>, <NAME>, <HERE>, <USING>, <GITHUB>, <HANDLE>
 
 
+# Release 2.9.2
+
+This patch release adds support for Tensorflow 2.16.
+
+## Bug Fixes and Other Changes
+
+* Support and test against Tensorflow 2.16. **Note that (like tensorflow-probability) GPflow uses Keras 2. Since TF 2.16 defaults to Keras 3, `tf.keras` must now be imported from the `tf_keras` package. Alternatively, you can import `tf_keras` from the `gpflow.keras` module, which will automatically select the right source depending on which version of TF is installed.**
+
+## Thanks to our Contributors
+
+This release contains contributions from:
+
+uri-granta
+
+
+# Release 2.9.1
+
+This patch release fixes a number of identified issues.
+
+## Bug Fixes and Other Changes
+
+* Support pickling Scipy optimizers with a non-empty compile cache
+* Allow setting a prior for A in the `Linear` mean function
+* Avoid rounding small values in kernel summary printout
+* Test against Tensorflow 2.15
+
+## Thanks to our Contributors
+
+This release contains contributions from:
+
+JSchmiegel, uri-granta
+
+
 # Release 2.9.0
 
 This release adds caching of compiled graphs inside the Scipy optimizer, and adds support

gpflow-2.9.2/VERSION

@@ -0,0 +1 @@
+2.9.2

{gpflow-2.9.0 → gpflow-2.9.2}/benchmark/plotters.py

@@ -18,8 +18,11 @@ from dataclasses import replace
 from math import isnan
 from typing import Collection, Tuple
 
+import matplotlib
 import pandas as pd
 from matplotlib.axes import Axes
+from matplotlib.dates import date2num
+from packaging.version import Version
 
 from benchmark.grouping import GroupingKey, GroupingSpec, group
 from benchmark.metadata import BenchmarkMetadata, parse_timestamp
@@ -84,8 +87,9 @@ def metrics_box_plot(
     for key, df, _ in group(metrics_df, [], metadata, line_by):
         labels.append(_join_key(key))
         values.append(df.value)
-    ax.boxplot(values, labels=labels)
-
+    # cf https://github.com/matplotlib/matplotlib/pull/27901
+    label_id = "tick_labels" if Version(matplotlib.__version__) >= Version("3.9.0") else "labels"
+    ax.boxplot(values, **{label_id: labels})  # type: ignore[arg-type]
     _shared_ax_config(ax, file_key, column_key, row_key, metric)
 
 
@@ -118,7 +122,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-2.9.0 → gpflow-2.9.2}/doc/sphinx/conf.py

@@ -24,7 +24,7 @@ copyright = "2023, The GPflow Contributors"
 author = "The GPflow Contributors"
 
 # The full version, including alpha/beta/rc tags
-release = "2.9.0"
+release = "2.9.2"
 
 # -- General configuration ---------------------------------------------------
 

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow/__init__.py

@@ -56,6 +56,7 @@ __all__ = [
     "experimental",
     "functions",
     "inducing_variables",
+    "keras",
     "kernels",
     "kullback_leiblers",
     "likelihoods",

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow/base.py

@@ -27,19 +27,19 @@ from .type_flags import GENERIC_NP_ARRAYS, NP_TYPE_CHECKING
 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-2.9.0 → gpflow-2.9.2}/gpflow/functions.py

@@ -110,7 +110,15 @@ class Linear(MeanFunction, Function):
         MeanFunction.__init__(self)
         A = np.ones((1, 1), dtype=default_float()) if A is None else A
         b = np.zeros(1, dtype=default_float()) if b is None else b
-        self.A = Parameter(np.atleast_2d(A))
+        if isinstance(A, Parameter):
+            if len(A._shape) >= 2:
+                self.A = A
+            else:
+                raise ValueError(
+                    "Error 'gpflow.funcitons.Linear()' mean function. A has not the correct shape (at least 2d)."
+                )
+        else:
+            self.A = Parameter(np.atleast_2d(A))
         self.b = Parameter(b)
 
     @inherit_check_shapes

gpflow-2.9.2/gpflow/keras.py

@@ -0,0 +1,32 @@
+# Copyright 2024 The GPflow Contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# GPflow currently uses Keras 2. Depending on the version of Tensorflow installed,
+# this is available either at tf_keras or at tensorflow.keras. This module provides
+# a helpful version-agnostic shortcut for importing this. Note though that importing specific
+# identifiers can only be done via attribute lookups:
+#
+# >> from gpflow.keras import tf_keras
+# >> Adam = tf_keras.optimizers.Adam
+# >> # The following does NOT work
+# >> from tf_keras.optimizers import Adam
+
+try:
+    import tf_keras
+except ModuleNotFoundError:
+    import tensorflow.keras as tf_keras
+
+tf_keras = tf_keras
+
+__all__ = ["tf_keras"]

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow/models/training_mixins.py

@@ -47,7 +47,7 @@ class InternalDataTrainingLossMixin:
       - a uniform API for the training loss :meth:`training_loss`
       - a convenience method :meth:`training_loss_closure` for constructing the closure expected by
         various optimizers, namely :class:`gpflow.optimizers.Scipy` and subclasses of
-        `tf.optimizers.Optimizer`.
+        `tf_keras.optimizers.Optimizer`.
 
     See :class:`ExternalDataTrainingLossMixin` for an equivalent mixin for models that do **not**
     own their own data.
@@ -67,7 +67,7 @@ class InternalDataTrainingLossMixin:
         """
         Convenience method. Returns a closure which itself returns the training loss. This closure
         can be passed to the minimize methods on :class:`gpflow.optimizers.Scipy` and subclasses of
-        `tf.optimizers.Optimizer`.
+        `tf_keras.optimizers.Optimizer`.
 
         :param compile: If `True` (default), compile the training loss function in a TensorFlow
             graph by wrapping it in tf.function()
@@ -86,7 +86,7 @@ class ExternalDataTrainingLossMixin:
       - a uniform API for the training loss :meth:`training_loss`
       - a convenience method :meth:`training_loss_closure` for constructing the closure expected by
         various optimizers, namely :class:`gpflow.optimizers.Scipy` and subclasses of
-        `tf.optimizers.Optimizer`.
+        `tf_keras.optimizers.Optimizer`.
 
     See :class:`InternalDataTrainingLossMixin` for an equivalent mixin for models that **do** own
     their own data.

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow/monitor/tensorboard.py

@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 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-2.9.0 → gpflow-2.9.2}/gpflow/optimizers/natgrad.py

@@ -16,13 +16,13 @@ import abc
 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
+from ..keras import tf_keras
 
-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"]]
 
@@ -173,7 +173,7 @@ class XiSqrtMeanVar(XiTransform):
         return natural_to_meanvarsqrt(nat1, nat2)
 
 
-class NaturalGradient(tf.optimizers.Optimizer):
+class NaturalGradient(tf_keras.optimizers.Optimizer):
     """
     Implements a natural gradient descent optimizer for variational models
     that are based on a distribution q(u) = N(q_mu, q_sqrt q_sqrtᵀ) that is
@@ -181,7 +181,7 @@ class NaturalGradient(tf.optimizers.Optimizer):
     of the covariance.
 
     Note that this optimizer does not implement the standard API of
-    tf.optimizers.Optimizer. Its only public method is minimize(), which has
+    tf_keras.optimizers.Optimizer. Its only public method is minimize(), which has
     a custom signature (var_list needs to be a list of (q_mu, q_sqrt) tuples,
     where q_mu and q_sqrt are gpflow.Parameter instances, not tf.Variable).
 

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow/optimizers/scipy.py

@@ -17,6 +17,7 @@ from collections import OrderedDict
 from typing import (
     Any,
     Callable,
+    Dict,
     FrozenSet,
     Iterable,
     List,
@@ -68,6 +69,12 @@ class Scipy:
             )
         self.compile_cache_size = compile_cache_size
 
+    def __getstate__(self) -> Dict[str, Any]:
+        # Don't try to save the compile cache
+        state = self.__dict__.copy()
+        state["compile_cache"] = OrderedDict()
+        return state
+
     def minimize(
         self,
         closure: LossClosure,

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow/utilities/misc.py

@@ -29,6 +29,8 @@ __all__ = [
     "training_loop",
 ]
 
+from ..keras import tf_keras
+
 
 @check_shapes(
     "x: [any...]",
@@ -73,7 +75,7 @@ def is_variable(t: TensorData) -> bool:
 
 def training_loop(
     closure: Callable[[], tf.Tensor],
-    optimizer: Optional[tf.optimizers.Optimizer] = None,
+    optimizer: Optional[tf_keras.optimizers.Optimizer] = None,
     var_list: Optional[List[tf.Variable]] = None,
     maxiter: int = 1_000,
     compile: bool = False,
@@ -90,7 +92,7 @@ def training_loop(
     :return:
     """
 
-    safe_optimizer = tf.optimizers.Adam() if optimizer is None else optimizer
+    safe_optimizer = tf_keras.optimizers.Adam() if optimizer is None else optimizer
     safe_var_list = [] if var_list is None else var_list
 
     def optimization_step() -> None:

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow/utilities/traversal.py

@@ -331,7 +331,7 @@ def _first_three_elements_regexp() -> Pattern[str]:
 
 
 def _str_tensor_value(value: AnyNDArray) -> str:
-    value_str = str(np.around(value, 5))
+    value_str = str(np.where(np.abs(value) < 1, value, np.around(value, 5)))
     if value.size <= 3:
         return value_str
 

{gpflow-2.9.0 → gpflow-2.9.2/gpflow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpflow
-Version: 2.9.0
+Version: 2.9.2
 Summary: Gaussian process methods in TensorFlow
 Home-page: https://www.gpflow.org
 Author: James Hensman, Alex Matthews
@@ -9,6 +9,7 @@ License: Apache License 2.0
 Project-URL: Source on GitHub, https://github.com/GPflow/GPflow
 Project-URL: Documentation, https://gpflow.github.io/GPflow/
 Keywords: machine-learning gaussian-processes kernels tensorflow
+Platform: UNKNOWN
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: GPU :: NVIDIA CUDA
 Classifier: Intended Audience :: Developers
@@ -216,3 +217,4 @@ which can be cited by users.
   url = {https://arxiv.org/abs/2003.01115}
 }
 ```
+

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow.egg-info/SOURCES.txt

@@ -34,6 +34,7 @@ gpflow/__init__.py
 gpflow/base.py
 gpflow/ci_utils.py
 gpflow/functions.py
+gpflow/keras.py
 gpflow/kullback_leiblers.py
 gpflow/logdensities.py
 gpflow/mean_functions.py

{gpflow-2.9.0 → gpflow-2.9.2}/gpflow.egg-info/requires.txt

@@ -6,7 +6,7 @@ packaging
 scipy
 setuptools>=41.0.0
 tabulate
-tensorflow-probability>=0.12.0
+tensorflow-probability[tf]>=0.12.0
 typing_extensions
 
 [:platform_system != "Darwin" or platform_machine != "arm64"]

{gpflow-2.9.0 → gpflow-2.9.2}/setup.py

@@ -16,7 +16,7 @@ requirements = [
     "scipy",
     "setuptools>=41.0.0",  # to satisfy dependency constraints
     "tabulate",
-    "tensorflow-probability>=0.12.0",
+    "tensorflow-probability[tf]>=0.12.0",
     "tensorflow>=2.4.0; platform_system!='Darwin' or platform_machine!='arm64'",
     # NOTE: Support of Apple Silicon MacOS platforms is in an experimental mode
     "tensorflow-macos>=2.4.0; platform_system=='Darwin' and platform_machine=='arm64'",

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/conditionals/test_uncertain_conditional.py

@@ -24,6 +24,7 @@ import gpflow
 from gpflow.base import AnyNDArray, MeanAndVariance
 from gpflow.conditionals import conditional, uncertain_conditional
 from gpflow.config import default_float
+from gpflow.keras import tf_keras
 from gpflow.mean_functions import Constant, Linear, MeanFunction, Zero
 from gpflow.quadrature import mvnquad
 from gpflow.utilities import training_loop
@@ -193,7 +194,7 @@ def test_no_uncertainty(white: bool, mean: Optional[str]) -> None:
 
     training_loop(
         model.training_loss_closure(Data.data),
-        optimizer=tf.optimizers.Adam(),
+        optimizer=tf_keras.optimizers.Adam(),
         var_list=model.trainable_variables,
         maxiter=100,
         compile=True,
@@ -226,7 +227,7 @@ def test_monte_carlo_1_din(white: bool, mean: Optional[str]) -> None:
 
     training_loop(
         model.training_loss_closure(DataMC1.data),
-        optimizer=tf.optimizers.Adam(),
+        optimizer=tf_keras.optimizers.Adam(),
         var_list=model.trainable_variables,
         maxiter=200,
         compile=True,
@@ -261,7 +262,7 @@ def test_monte_carlo_2_din(white: bool, mean: Optional[str]) -> None:
 
     training_loop(
         model.training_loss_closure(DataMC2.data),
-        optimizer=tf.optimizers.Adam(),
+        optimizer=tf_keras.optimizers.Adam(),
         var_list=model.trainable_variables,
         maxiter=100,
         compile=True,

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/likelihoods/test_multiclass.py

@@ -48,7 +48,7 @@ def test_softmax_y_shape_assert(num: int, dimX: int, dimF: int, dimY: int) -> No
 
 
 @pytest.mark.parametrize("dimX", [3])
-@pytest.mark.parametrize("num", [10, 3])
+@pytest.mark.parametrize("num", [8, 3])
 @pytest.mark.parametrize("dimF, dimY", [[2, 1]])
 def test_softmax_bernoulli_equivalence(num: int, dimX: int, dimF: int, dimY: int) -> None:
     X = tf.random.normal((num, dimX))

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/models/test_svgp.py

@@ -21,6 +21,7 @@ from numpy.testing import assert_allclose
 import gpflow
 from gpflow import set_trainable
 from gpflow.base import AnyNDArray
+from gpflow.keras import tf_keras
 from gpflow.models import SVGP
 
 
@@ -184,7 +185,7 @@ def test_stochastic_gradients(
 
     def training_loop(indices: Sequence[int], num_data: int, max_iter: int) -> SVGP:
         model = get_model(num_data)
-        opt = tf.optimizers.SGD(learning_rate=0.001)
+        opt = tf_keras.optimizers.SGD(learning_rate=0.001)
         data = X[indices], Y[indices]
         for _ in range(max_iter):
             with tf.GradientTape() as tape:

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/optimizers/test_scipy.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import copy
 import unittest
 import warnings
 from typing import Any, Callable, Dict, List, Sequence, Tuple
@@ -545,3 +545,35 @@ def test_scipy__cache_with_different_models() -> None:
 
     # Check that the cache was bypassed and a new compiled function was created
     assert len(opt.compile_cache) == 2
+
+
+def test_scipy__deep_copyable() -> None:
+    opt = gpflow.optimizers.Scipy()
+    variables1 = _create_variables()
+    counter1 = [0]
+    closure1 = _loss_closure_builder(counter1, variables1)
+
+    # Ensure the cache is empty
+    assert len(opt.compile_cache) == 0
+
+    # Call eval_func multiple times with the same arguments
+    for _ in range(4):
+        opt.minimize(
+            closure1, variables1, compile=True, tf_fun_args={}, allow_unused_variables=False
+        )
+        # Check that the cache was used
+        assert len(opt.compile_cache) == 1
+
+    # Check that the optimizer can still be deepcopied
+    opt_copy = copy.deepcopy(opt)
+
+    # The cache is not copied
+    assert len(opt_copy.compile_cache) == 0
+
+    # Call eval_func multiple times with the same arguments
+    for _ in range(4):
+        opt_copy.minimize(
+            closure1, variables1, compile=True, tf_fun_args={}, allow_unused_variables=False
+        )
+        # Check that the cache was used
+        assert len(opt_copy.compile_cache) == 1

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/test_all.py

@@ -58,14 +58,11 @@ _MODULES_WITH_ALL = [m for m in _MODULES if hasattr(m, "__all__")]
 
 @pytest.mark.parametrize("package", _PACKAGES)
 def test_all_present_and_up_to_date(package: ModuleType) -> None:
-    imported_sorted = sorted(attr for attr in dir(package) if not is_dunder(attr))
+    imported_set = set(attr for attr in dir(package) if not is_dunder(attr))
     all_list = getattr(package, "__all__", None)
     assert all_list is not None, f"Package {package} is missing an explicit __all__."
-    all_sorted = sorted(attr for attr in all_list if not is_dunder(attr))
-    assert imported_sorted == all_sorted, (
-        f"{package}.__all__ is outdated."
-        f" Imported values are {imported_sorted}, but exported values are {all_list}."
-    )
+    all_set = set(attr for attr in all_list if not is_dunder(attr))
+    assert imported_set == all_set, f"{package}.__all__ is outdated."
 
 
 @pytest.mark.parametrize("module", _MODULES_WITH_ALL)

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/test_functions.py

@@ -16,6 +16,7 @@ from typing import Any, Sequence, Tuple, Type
 
 import numpy as np
 import pytest
+import tensorflow_probability as tfp
 from check_shapes import check_shapes
 from numpy.testing import assert_allclose
 
@@ -371,3 +372,25 @@ def test_models_with_mean_functions_changes(model_class: Type[Any]) -> None:
     assert np.all(var_zero.numpy() == var_non_zero.numpy())
     # predictive mean changes after modifying mean function
     assert not np.all(mu_zero.numpy() == mu_non_zero.numpy())
+
+
+class TestIssue2091EnsureParameterTypeForLinearMeanFunction:
+    """
+    See github issue #2091. These are tests to ensure that the prior is kept if A is already given as Parameter.
+    """
+
+    # test_parameter_with_correct_shape
+    def test_parameter_with_correct_shape(self) -> None:
+        "Check that Linear copies the prior from a correctly shaped A Parameter"
+        A = gpflow.Parameter(
+            np.ones((1, 1)), dtype=np.float64, prior=tfp.distributions.Normal(0.0, 1.0)
+        )
+        linear_function = Linear(A, 1)
+        assert linear_function.A.prior is not None
+
+    # test_parameter_with_incorrect_shape
+    def test_parameter_with_incorrect_shape(self) -> None:
+        "Check that Linear throws an error when the A Parameter is not correctly shaped"
+        A = gpflow.Parameter(np.zeros(20), dtype=np.float64)
+        with pytest.raises(ValueError):
+            Linear(A, 1)

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/test_inducing_variables.py

@@ -23,6 +23,7 @@ from check_shapes.error_contexts import ErrorContext, MessageBuilder
 
 import gpflow
 import gpflow.inducing_variables as giv
+from gpflow.keras import tf_keras
 
 
 @dataclass(frozen=True)
@@ -50,7 +51,7 @@ def test_inducing_points_with_variable_shape() -> None:
     m = gpflow.models.SGPR(data=(X, Y), kernel=gpflow.kernels.Matern32(), inducing_variable=iv)
 
     # Check 1: that we can still optimize with None shape
-    opt = tf.optimizers.Adam()
+    opt = tf_keras.optimizers.Adam()
 
     @tf.function
     def optimization_step() -> None:

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/test_monitor.py

@@ -11,6 +11,8 @@ from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 
 import gpflow
+from gpflow.base import AnyNDArray
+from gpflow.keras import tf_keras
 from gpflow.models import GPR, GPModel
 from gpflow.monitor import (
     ExecuteCallback,
@@ -105,7 +107,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))
@@ -303,7 +305,7 @@ def test_logdir_created(monitor: Monitor, model: GPModel, tmp_path: Path) -> Non
     size_before = _get_size_directory(tmp_path)
     assert size_before > 0
 
-    opt = tf.optimizers.Adam()
+    opt = tf_keras.optimizers.Adam()
     for step in range(Data.num_steps):
         opt.minimize(model.training_loss, model.trainable_variables)
         monitor(step)
@@ -313,7 +315,7 @@ def test_logdir_created(monitor: Monitor, model: GPModel, tmp_path: Path) -> Non
 
 
 def test_compile_monitor(monitor: Monitor, model: GPModel) -> None:
-    opt = tf.optimizers.Adam()
+    opt = tf_keras.optimizers.Adam()
 
     @tf.function
     def tf_func(step: tf.Tensor) -> None:

{gpflow-2.9.0 → gpflow-2.9.2}/tests/gpflow/utilities/test_training_loop.py

@@ -1,10 +1,10 @@
 from typing import Any
 
 import numpy as np
-import tensorflow as tf
 from check_shapes import ShapeChecker
 
 import gpflow
+from gpflow.keras import tf_keras
 from gpflow.utilities import training_loop
 
 
@@ -43,10 +43,18 @@ def test_training_loop_compiles() -> None:
     m1 = create_model()
     m2 = create_model()
     training_loop(
-        m1.training_loss, tf.optimizers.Adam(), m1.trainable_variables, maxiter=50, compile=True
+        m1.training_loss,
+        tf_keras.optimizers.Adam(),
+        m1.trainable_variables,
+        maxiter=50,
+        compile=True,
     )
     training_loop(
-        m2.training_loss, tf.optimizers.Adam(), m2.trainable_variables, maxiter=50, compile=False
+        m2.training_loss,
+        tf_keras.optimizers.Adam(),
+        m2.trainable_variables,
+        maxiter=50,
+        compile=False,
     )
     assert_models_close(m1, m2)
 
@@ -57,7 +65,7 @@ def test_training_loop_converges() -> None:
     gpflow.optimizers.Scipy().minimize(mref.training_loss, mref.trainable_variables)
     training_loop(
         m.training_loss,
-        tf.optimizers.Adam(learning_rate=0.01),
+        tf_keras.optimizers.Adam(learning_rate=0.01),
         m.trainable_variables,
         maxiter=5000,
         compile=True,