Support tensorflow 2.16 (#102)

.github/workflows/quality-check.yaml

@@ -20,9 +20,10 @@ jobs:
   check-and-test:
     runs-on: ubuntu-20.04
     strategy:
+      fail-fast: false
       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 +58,8 @@ jobs:
             - tensorflow: "~=2.15.0"
               keras: "~=2.15.0"
               tensorflow-probability: "~=0.23.0"
+            - tensorflow: "~=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

@@ -27,6 +27,8 @@
 from scipy.stats import norm
 from utils import ExperimentName, git_version
 
+from gpflow.keras import tf_keras
+
 from gpflux.architectures import Config, build_constant_input_dim_deep_gp
 
 THIS_DIR = Path(__file__).parent
@@ -89,10 +91,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
         ),
     ]
@@ -123,7 +125,7 @@ def main(_config):
     data = get_data()
     model = build_model(data.X_train)
 
-    model.compile(optimizer=tf.optimizers.Adam(0.01))
+    model.compile(optimizer=tf_keras.optimizers.Adam(0.01))
     train_model(model, (data.X_train, data.Y_train))
 
     metrics = evaluate_model(model, (data.X_test, data.Y_test))

docs/index.rst

@@ -71,7 +71,7 @@ As a quick teaser, here's a snippet from the `intro notebook <notebooks/intro>`
 
    # Compile and fit
    model = two_layer_dgp.as_training_model()
-   model.compile(tf.optimizers.Adam(0.01))
+   model.compile(gpflow.keras.tf_keras.optimizers.Adam(0.01))
    history = model.fit({"inputs": X, "targets": Y}, epochs=int(1e3), verbose=0)
 
 The model described above produces the fit shown in Fig 1. For comparison, in Fig. 2 we show the fit on the same dataset by a vanilla single-layer GP model.

docs/notebooks/deep_cde.ipynb

@@ -124,7 +124,7 @@
     "\n",
     "single_layer_dgp = gpflux.models.DeepGP([gp_layer], likelihood_layer)\n",
     "model = single_layer_dgp.as_training_model()\n",
-    "model.compile(tf.optimizers.Adam(0.01))\n",
+    "model.compile(gpflow.keras.tf_keras.optimizers.Adam(0.01))\n",
     "\n",
     "history = model.fit({\"inputs\": X, \"targets\": Y}, epochs=int(1e3), verbose=0)\n",
     "fig, ax = plt.subplots()\n",
@@ -387,7 +387,7 @@
    "execution_count": 17,
    "source": [
     "model = dgp.as_training_model()\n",
-    "model.compile(tf.optimizers.Adam(0.005))\n",
+    "model.compile(gpflow.keras.tf_keras.optimizers.Adam(0.005))\n",
     "history = model.fit({\"inputs\": X, \"targets\": Y}, epochs=int(20e3), verbose=0, batch_size=num_data, shuffle=False)"
    ],
    "outputs": [],

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
@@ -99,10 +100,10 @@
 """
 
 # %%
-model.compile(tf.optimizers.Adam(learning_rate=0.1))
+model.compile(tf_keras.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))
+training_model.compile(optimizer=tf_keras.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/intro.py

@@ -97,7 +97,7 @@ def motorcycle_data():
 # %%
 single_layer_dgp = gpflux.models.DeepGP([gp_layer], likelihood_layer)
 model = single_layer_dgp.as_training_model()
-model.compile(tf.optimizers.Adam(0.01))
+model.compile(gpflow.keras.tf_keras.optimizers.Adam(0.01))
 
 # %% [markdown]
 """
@@ -168,7 +168,7 @@ def plot(model, X, Y, ax=None):
 likelihood_layer = gpflux.layers.LikelihoodLayer(gpflow.likelihoods.Gaussian(0.1))
 two_layer_dgp = gpflux.models.DeepGP([gp_layer1, gp_layer2], likelihood_layer)
 model = two_layer_dgp.as_training_model()
-model.compile(tf.optimizers.Adam(0.01))
+model.compile(gpflow.keras.tf_keras.optimizers.Adam(0.01))
 
 # %%
 history = model.fit({"inputs": X, "targets": Y}, epochs=int(1e3), verbose=0)

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,
@@ -96,7 +97,7 @@ def create_model(model_class):
 ]
 
 dgp_train = dgp.as_training_model()
-dgp_train.compile(tf.optimizers.Adam(learning_rate=0.1))
+dgp_train.compile(tf_keras.optimizers.Adam(learning_rate=0.1))
 
 history = dgp_train.fit(
     {"inputs": X, "targets": Y}, batch_size=batch_size, epochs=num_epochs, callbacks=callbacks
@@ -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,
@@ -124,7 +125,7 @@ def create_model(model_class):
     [
         gpflow.optimizers.NaturalGradient(gamma=0.05),
         gpflow.optimizers.NaturalGradient(gamma=0.05),
-        tf.optimizers.Adam(learning_rate=0.1),
+        tf_keras.optimizers.Adam(learning_rate=0.1),
     ]
 )
 

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/helpers.py

@@ -22,7 +22,7 @@
 import inspect
 import warnings
 from dataclasses import fields
-from typing import List, Optional, Type, TypeVar, Union, Any
+from typing import Any, List, Optional, Type, TypeVar, Union
 
 import numpy as np
 

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,9 +15,10 @@
 #
 """Utility layer that tracks variables in :class:`tf.Module`."""
 
-import tensorflow as tf
 from deprecated import deprecated
 
+from gpflow.keras import tf_keras
+
 
 @deprecated(
     reason=(
@@ -27,7 +28,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.