Support tensorflow 2.16

.github/workflows/quality-check.yaml

@@ -22,7 +22,7 @@ jobs:
     strategy:
       matrix:
           python-version: ["3.7", "3.8", "3.9", "3.10"]
-          tensorflow: ["~=2.5.0", "~=2.6.0", "~=2.7.0", "~=2.8.0", "~=2.9.0", "~=2.10.0", "~=2.11.0", "~=2.12.0", "~=2.13.0", "~=2.14.0", "~=2.15.0"]
+          tensorflow: ["~=2.5.0", "~=2.6.0", "~=2.7.0", "~=2.8.0", "~=2.9.0", "~=2.10.0", "~=2.11.0", "~=2.12.0", "~=2.13.0", "~=2.14.0", "~=2.15.0", "~=2.16.0"]
           include:
             - tensorflow: "~=2.5.0"
               keras: "~=2.6.0"
@@ -57,6 +57,9 @@ jobs:
             - tensorflow: "~=2.15.0"
               keras: "~=2.15.0"
               tensorflow-probability: "~=0.23.0"
+            - tensorflow: "~=2.16.0"
+              keras: "~=2.16.0"
+              tensorflow-probability: "~=0.24.0"
           exclude:
             # These older versions of TensorFlow don't work with Python 3.10:
             - python-version: "3.10"
@@ -74,11 +77,15 @@ jobs:
               tensorflow: "~=2.14.0"
             - python-version: "3.7"
               tensorflow: "~=2.15.0"
+            - python-version: "3.7"
+              tensorflow: "~=2.16.0"
             # These newer versions of TensorFlow don't work with Python 3.8:
             - python-version: "3.8"
               tensorflow: "~=2.14.0"
             - python-version: "3.8"
               tensorflow: "~=2.15.0"
+            - python-version: "3.8"
+              tensorflow: "~=2.16.0"
 
     name: Python-${{ matrix.python-version }} tensorflow${{ matrix.tensorflow }}
     env:

benchmarking/main.py

@@ -23,6 +23,7 @@
 import tensorflow as tf
 from bayesian_benchmarks import data as uci_datasets
 from bayesian_benchmarks.data import Dataset
+from gpflow.keras import tf_keras
 from sacred import Experiment
 from scipy.stats import norm
 from utils import ExperimentName, git_version
@@ -89,10 +90,10 @@ def build_model(X, num_inducing, num_layers):
 
 
 @EXPERIMENT.capture
-def train_model(model: tf.keras.models.Model, data_train, batch_size, num_epochs):
+def train_model(model: tf_keras.models.Model, data_train, batch_size, num_epochs):
     X_train, Y_train = data_train
     callbacks = [
-        tf.keras.callbacks.ReduceLROnPlateau(
+        tf_keras.callbacks.ReduceLROnPlateau(
             'loss', factor=0.95, patience=3, min_lr=1e-6, verbose=1
         ),
     ]

docs/notebooks/efficient_sampling.py

@@ -35,6 +35,7 @@
 import gpflux
 
 from gpflow.config import default_float
+from gpflow.keras import tf_keras
 
 from gpflux.layers.basis_functions.fourier_features import RandomFourierFeaturesCosine
 from gpflux.sampling import KernelWithFeatureDecomposition
@@ -102,7 +103,7 @@
 model.compile(tf.optimizers.Adam(learning_rate=0.1))
 
 callbacks = [
-    tf.keras.callbacks.ReduceLROnPlateau(
+    tf_keras.callbacks.ReduceLROnPlateau(
         monitor="loss",
         patience=5,
         factor=0.95,

docs/notebooks/gpflux_features.py

@@ -63,6 +63,7 @@ def motorcycle_data():
 # %%
 import gpflux
 
+from gpflow.keras import tf_keras
 from gpflux.architectures import Config, build_constant_input_dim_deep_gp
 from gpflux.models import DeepGP
 
@@ -80,19 +81,19 @@ def motorcycle_data():
 
 # %%
 # From the `DeepGP` model we instantiate a training model which is a `tf.keras.Model`
-training_model: tf.keras.Model = deep_gp.as_training_model()
+training_model: tf_keras.Model = deep_gp.as_training_model()
 
 # Following the Keras procedure we need to compile and pass a optimizer,
 # before fitting the model to data
 training_model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01))
 
 callbacks = [
     # Create callback that reduces the learning rate every time the ELBO plateaus
-    tf.keras.callbacks.ReduceLROnPlateau("loss", factor=0.95, patience=3, min_lr=1e-6, verbose=0),
+    tf_keras.callbacks.ReduceLROnPlateau("loss", factor=0.95, patience=3, min_lr=1e-6, verbose=0),
     # Create a callback that writes logs (e.g., hyperparameters, KLs, etc.) to TensorBoard
     gpflux.callbacks.TensorBoard(),
     # Create a callback that saves the model's weights
-    tf.keras.callbacks.ModelCheckpoint(filepath="ckpts/", save_weights_only=True, verbose=0),
+    tf_keras.callbacks.ModelCheckpoint(filepath="ckpts/", save_weights_only=True, verbose=0),
 ]
 
 history = training_model.fit(

docs/notebooks/gpflux_with_keras_layers.py

@@ -29,6 +29,7 @@
 import gpflux
 
 from gpflow.config import default_float
+from gpflow.keras import tf_keras
 
 
 # %% [markdown]
@@ -78,11 +79,11 @@
 likelihood_container = gpflux.layers.TrackableLayer()
 likelihood_container.likelihood = likelihood
 
-model = tf.keras.Sequential(
+model = tf_keras.Sequential(
     [
-        tf.keras.layers.Dense(100, activation="relu"),
-        tf.keras.layers.Dense(100, activation="relu"),
-        tf.keras.layers.Dense(1, activation="linear"),
+        tf_keras.layers.Dense(100, activation="relu"),
+        tf_keras.layers.Dense(100, activation="relu"),
+        tf_keras.layers.Dense(1, activation="linear"),
         gp_layer,
         likelihood_container,  # no-op, for discovering trainable likelihood parameters
     ]

docs/notebooks/keras_integration.py

@@ -25,6 +25,7 @@
 import gpflow
 import gpflux
 from gpflow.ci_utils import reduce_in_tests
+from gpflow.keras import tf_keras
 
 import matplotlib.pyplot as plt
 
@@ -83,10 +84,10 @@ def create_model(model_class):
 num_epochs = reduce_in_tests(200)
 
 # %%
-dgp = create_model(tf.keras.Model)
+dgp = create_model(tf_keras.Model)
 
 callbacks = [
-    tf.keras.callbacks.ReduceLROnPlateau(
+    tf_keras.callbacks.ReduceLROnPlateau(
         monitor="loss",
         patience=5,
         factor=0.95,
@@ -106,7 +107,7 @@ def create_model(model_class):
 dgp_natgrad = create_model(gpflux.optimization.NatGradModel)
 
 callbacks = [
-    tf.keras.callbacks.ReduceLROnPlateau(
+    tf_keras.callbacks.ReduceLROnPlateau(
         monitor="loss",
         patience=5,
         factor=0.95,

gpflux/callbacks.py

@@ -21,15 +21,14 @@
 import re
 from typing import Any, Dict, List, Mapping, Optional, Union
 
-import tensorflow as tf
-
 import gpflow
+from gpflow.keras import tf_keras
 from gpflow.utilities import parameter_dict
 
 __all__ = ["TensorBoard"]
 
 
-class TensorBoard(tf.keras.callbacks.TensorBoard):
+class TensorBoard(tf_keras.callbacks.TensorBoard):
     """
     This class is a thin wrapper around a `tf.keras.callbacks.TensorBoard` callback that also
     calls GPflow's `gpflow.monitor.ModelToTensorBoard` monitoring task.
@@ -100,7 +99,7 @@ def __init__(
         self.keywords_to_monitor = keywords_to_monitor
         self.max_size = max_size
 
-    def set_model(self, model: tf.keras.Model) -> None:
+    def set_model(self, model: tf_keras.Model) -> None:
         """
         Set the model (extends the Keras `set_model
         <https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/TensorBoard#set_model>`_

gpflux/layers/basis_functions/fourier_features/base.py

@@ -22,11 +22,12 @@
 
 import gpflow
 from gpflow.base import TensorType
+from gpflow.keras import tf_keras
 
 from gpflux.types import ShapeType
 
 
-class FourierFeaturesBase(ABC, tf.keras.layers.Layer):
+class FourierFeaturesBase(ABC, tf_keras.layers.Layer):
     r"""
     The base class for all Fourier feature layers, used for both random Fourier feature layers and
     quadrature layers. We subclass :class:`tf.keras.layers.Layer`, so we must provide

gpflux/layers/latent_variable_layer.py

@@ -23,6 +23,7 @@
 
 from gpflow import default_float
 from gpflow.base import TensorType
+from gpflow.keras import tf_keras
 
 from gpflux.layers.trackable_layer import TrackableLayer
 from gpflux.types import ObservationType
@@ -67,14 +68,14 @@ class LatentVariableLayer(LayerWithObservations):
     prior: tfp.distributions.Distribution
     """ The prior distribution for the latent variables. """
 
-    encoder: tf.keras.layers.Layer
+    encoder: tf_keras.layers.Layer
     """
     An encoder that maps from a concatenation of inputs and targets to the
     parameters of the approximate posterior distribution of the corresponding
     latent variables.
     """
 
-    compositor: tf.keras.layers.Layer
+    compositor: tf_keras.layers.Layer
     """
     A layer that takes as input the two-element ``[layer_inputs, latent_variable_samples]`` list
     and combines the elements into a single output tensor.
@@ -83,8 +84,8 @@ class LatentVariableLayer(LayerWithObservations):
     def __init__(
         self,
         prior: tfp.distributions.Distribution,
-        encoder: tf.keras.layers.Layer,
-        compositor: Optional[tf.keras.layers.Layer] = None,
+        encoder: tf_keras.layers.Layer,
+        compositor: Optional[tf_keras.layers.Layer] = None,
         name: Optional[str] = None,
     ):
         """
@@ -108,7 +109,7 @@ def __init__(
         self.compositor = (
             compositor
             if compositor is not None
-            else tf.keras.layers.Concatenate(axis=-1, dtype=default_float())
+            else tf_keras.layers.Concatenate(axis=-1, dtype=default_float())
         )
 
     def call(

gpflux/layers/trackable_layer.py

@@ -15,8 +15,8 @@
 #
 """Utility layer that tracks variables in :class:`tf.Module`."""
 
-import tensorflow as tf
 from deprecated import deprecated
+from gpflow.keras import tf_keras
 
 
 @deprecated(
@@ -27,7 +27,7 @@
         "be removed in GPflux version `1.0.0`."
     )
 )
-class TrackableLayer(tf.keras.layers.Layer):
+class TrackableLayer(tf_keras.layers.Layer):
     """
     With the release of TensorFlow 2.5, our TrackableLayer workaround is no
     longer needed.  See https://github.com/Prowler-io/gpflux/issues/189.

gpflux/losses.py

@@ -25,11 +25,12 @@
 
 import gpflow
 from gpflow.base import TensorType
+from gpflow.keras import tf_keras
 
 from gpflux.types import unwrap_dist
 
 
-class LikelihoodLoss(tf.keras.losses.Loss):
+class LikelihoodLoss(tf_keras.losses.Loss):
     r"""
     This class is a `tf.keras.losses.Loss` implementation that wraps a GPflow
     :class:`~gpflow.likelihoods.Likelihood` instance.

gpflux/models/deep_gp.py

@@ -22,6 +22,7 @@
 
 import gpflow
 from gpflow.base import Module, TensorType
+from gpflow.keras import tf_keras
 
 import gpflux
 from gpflux.layers import LayerWithObservations, LikelihoodLayer
@@ -46,16 +47,16 @@ class DeepGP(Module):
        for more details.
     """
 
-    inputs: tf.keras.Input
-    targets: tf.keras.Input
+    inputs: tf_keras.Input
+    targets: tf_keras.Input
 
-    f_layers: List[tf.keras.layers.Layer]
+    f_layers: List[tf_keras.layers.Layer]
     """ A list of all layers in this DeepGP (just :attr:`likelihood_layer` is separate). """
 
     likelihood_layer: gpflux.layers.LikelihoodLayer
     """ The likelihood layer. """
 
-    default_model_class: Type[tf.keras.Model]
+    default_model_class: Type[tf_keras.Model]
     """
     The default for the *model_class* argument of :meth:`as_training_model` and
     :meth:`as_prediction_model`. This must have the same semantics as `tf.keras.Model`,
@@ -73,14 +74,14 @@ class DeepGP(Module):
 
     def __init__(
         self,
-        f_layers: List[tf.keras.layers.Layer],
+        f_layers: List[tf_keras.layers.Layer],
         likelihood: Union[
             gpflux.layers.LikelihoodLayer, gpflow.likelihoods.Likelihood
         ],  # fully-qualified for autoapi
         *,
         input_dim: Optional[int] = None,
         target_dim: Optional[int] = None,
-        default_model_class: Type[tf.keras.Model] = tf.keras.Model,
+        default_model_class: Type[tf_keras.Model] = tf_keras.Model,
         num_data: Optional[int] = None,
     ):
         """
@@ -99,8 +100,8 @@ def __init__(
             If you do not specify a value for this parameter explicitly, it is automatically
             detected from the :attr:`~gpflux.layers.GPLayer.num_data` attribute in the GP layers.
         """
-        self.inputs = tf.keras.Input((input_dim,), dtype=gpflow.default_float(), name="inputs")
-        self.targets = tf.keras.Input((target_dim,), dtype=gpflow.default_float(), name="targets")
+        self.inputs = tf_keras.Input((input_dim,), dtype=gpflow.default_float(), name="inputs")
+        self.targets = tf_keras.Input((target_dim,), dtype=gpflow.default_float(), name="targets")
         self.f_layers = f_layers
         if isinstance(likelihood, gpflow.likelihoods.Likelihood):
             self.likelihood_layer = LikelihoodLayer(likelihood)
@@ -111,7 +112,7 @@ def __init__(
 
     @staticmethod
     def _validate_num_data(
-        f_layers: List[tf.keras.layers.Layer], num_data: Optional[int] = None
+        f_layers: List[tf_keras.layers.Layer], num_data: Optional[int] = None
     ) -> int:
         """
         Check that the :attr:`~gpflux.layers.gp_layer.GPLayer.num_data`
@@ -229,15 +230,15 @@ def elbo(self, data: Tuple[TensorType, TensorType]) -> tf.Tensor:
         ]
         return -tf.reduce_sum(all_losses) * self.num_data
 
-    def _get_model_class(self, model_class: Optional[Type[tf.keras.Model]]) -> Type[tf.keras.Model]:
+    def _get_model_class(self, model_class: Optional[Type[tf_keras.Model]]) -> Type[tf_keras.Model]:
         if model_class is not None:
             return model_class
         else:
             return self.default_model_class
 
     def as_training_model(
-        self, model_class: Optional[Type[tf.keras.Model]] = None
-    ) -> tf.keras.Model:
+        self, model_class: Optional[Type[tf_keras.Model]] = None
+    ) -> tf_keras.Model:
         r"""
         Construct a `tf.keras.Model` instance that requires you to provide both ``inputs``
         and ``targets`` to its call. This information is required for
@@ -269,8 +270,8 @@ def as_training_model(
         return model_class([self.inputs, self.targets], outputs)
 
     def as_prediction_model(
-        self, model_class: Optional[Type[tf.keras.Model]] = None
-    ) -> tf.keras.Model:
+        self, model_class: Optional[Type[tf_keras.Model]] = None
+    ) -> tf_keras.Model:
         """
         Construct a `tf.keras.Model` instance that requires only ``inputs``,
         which means you do not have to provide dummy target values when