removing dependency on setting the Keras backend

benchmarking/main.py

@@ -29,8 +29,6 @@
 
 from gpflux.architectures import Config, build_constant_input_dim_deep_gp
 
-tf.keras.backend.set_floatx("float64")
-
 THIS_DIR = Path(__file__).parent
 LOGS = THIS_DIR / "tmp"
 EXPERIMENT = Experiment("UCI")

docs/notebooks/deep_cde.ipynb

@@ -23,10 +23,7 @@
     "from tqdm import tqdm\n",
     "\n",
     "import tensorflow_probability as tfp\n",
-    "from sklearn.neighbors import KernelDensity\n",
-    "\n",
-    "\n",
-    "tf.keras.backend.set_floatx(\"float64\")"
+    "from sklearn.neighbors import KernelDensity\n"
    ],
    "outputs": [],
    "metadata": {}

docs/notebooks/efficient_sampling.py

@@ -40,7 +40,6 @@
 from gpflux.sampling import KernelWithFeatureDecomposition
 from gpflux.models.deep_gp import sample_dgp
 
-tf.keras.backend.set_floatx("float64")
 
 # %% [markdown]
 """

docs/notebooks/gpflux_features.py

@@ -33,7 +33,6 @@
 import pandas as pd
 import tensorflow as tf
 
-tf.keras.backend.set_floatx("float64")  # we want to carry out GP calculations in 64 bit
 tf.get_logger().setLevel("INFO")
 
 

docs/notebooks/gpflux_with_keras_layers.py

@@ -30,7 +30,6 @@
 
 from gpflow.config import default_float
 
-tf.keras.backend.set_floatx("float64")
 
 # %% [markdown]
 """

docs/notebooks/intro.py

@@ -26,7 +26,6 @@
 import pandas as pd
 import tensorflow as tf
 
-tf.keras.backend.set_floatx("float64")
 tf.get_logger().setLevel("INFO")
 
 # %% [markdown]

docs/notebooks/keras_integration.py

@@ -28,8 +28,6 @@
 
 import matplotlib.pyplot as plt
 
-# %%
-tf.keras.backend.set_floatx("float64")
 
 # %%
 # %matplotlib inline

gpflux/architectures/constant_input_dim_deep_gp.py

@@ -113,6 +113,12 @@ def build_constant_input_dim_deep_gp(X: np.ndarray, num_layers: int, config: Con
     :param num_layers: The number of layers in the Deep GP.
     :param config: The configuration for (hyper)parameters. See :class:`Config` for details.
     """
+    if X.dtype != gpflow.default_float():
+        raise ValueError(
+            f"X needs to have dtype according to gpflow.default_float() = {gpflow.default_float()} "
+            f"however got X with {X.dtype} dtype."
+        )
+
     num_data, input_dim = X.shape
     X_running = X
 

gpflux/layers/latent_variable_layer.py

@@ -94,7 +94,10 @@ def __init__(
             posterior distribution; see :attr:`encoder`.
         :param compositor: A layer that combines layer inputs and latent variable
             samples into a single tensor; see :attr:`compositor`. If you do not specify a value for
-            this parameter, the default is ``tf.keras.layers.Concatenate(axis=-1)``.
+            this parameter, the default is
+            ``tf.keras.layers.Concatenate(axis=-1, dtype=default_float())``. Note that you should
+            set ``dtype`` of the layer to GPflow's default dtype as in
+            :meth:`~gpflow.default_float()`.
         :param name: The name of this layer (passed through to `tf.keras.layers.Layer`).
         """
 
@@ -103,7 +106,9 @@ def __init__(
         self.distribution_class = prior.__class__
         self.encoder = encoder
         self.compositor = (
-            compositor if compositor is not None else tf.keras.layers.Concatenate(axis=-1)
+            compositor
+            if compositor is not None
+            else tf.keras.layers.Concatenate(axis=-1, dtype=default_float())
         )
 
     def call(

gpflux/models/deep_gp.py

@@ -37,6 +37,9 @@ class DeepGP(Module):
     inheriting from :class:`~gpflux.layers.LayerWithObservations`; those will
     be passed the argument ``observations=[inputs, targets]``.
 
+    When data is used with methods in this class (e.g. :meth:`predict_f` method), it needs to
+    be with ``dtype`` corresponding to GPflow's default dtype as in :meth:`~gpflow.default_float()`.
+
     .. note:: This class is **not** a `tf.keras.Model` subclass itself. To access
        Keras features, call either :meth:`as_training_model` or :meth:`as_prediction_model`
        (depending on the use-case) to create a `tf.keras.Model` instance. See the method docstrings
@@ -96,8 +99,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,), name="inputs")
-        self.targets = tf.keras.Input((target_dim,), 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)
@@ -130,6 +133,20 @@ def _validate_num_data(
             raise ValueError("Could not determine num_data; please provide explicitly")
         return num_data
 
+    @staticmethod
+    def _validate_dtype(x: TensorType) -> None:
+        """
+        Check that data ``x`` is of correct ``dtype``, corresponding to GPflow's default dtype as
+        defined by :meth:`~gpflow.default_float()`.
+
+        :raise ValueError: If ``x`` is of incorrect ``dtype.
+        """
+        if x.dtype != gpflow.default_float():
+            raise ValueError(
+                f"x needs to have dtype {gpflow.default_float()} (according to "
+                f"gpflow.default_float()), however got x with {x.dtype} dtype."
+            )
+
     def _evaluate_deep_gp(
         self,
         inputs: TensorType,
@@ -180,6 +197,9 @@ def call(
         targets: Optional[TensorType] = None,
         training: Optional[bool] = None,
     ) -> tf.Tensor:
+        self._validate_dtype(inputs)
+        if targets is not None:
+            self._validate_dtype(targets)
         f_outputs = self._evaluate_deep_gp(inputs, targets=targets, training=training)
         y_outputs = self._evaluate_likelihood(f_outputs, targets=targets, training=training)
         return y_outputs
@@ -188,9 +208,11 @@ def predict_f(self, inputs: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
         """
         :returns: The mean and variance (not the scale!) of ``f``, for compatibility with GPflow
            models.
+        :raise ValueError: If ``x`` is of incorrect ``dtype.
 
         .. note:: This method does **not** support ``full_cov`` or ``full_output_cov``.
         """
+        self._validate_dtype(inputs)
         f_distribution = self._evaluate_deep_gp(inputs, targets=None)
         return f_distribution.loc, f_distribution.scale.diag ** 2
 

tests/gpflux/architectures/test_constant_input_dim_deep_gp.py

@@ -5,8 +5,6 @@
 from gpflux.architectures import Config, build_constant_input_dim_deep_gp
 from gpflux.helpers import make_dataclass_from_class
 
-tf.keras.backend.set_floatx("float64")
-
 
 class DemoConfig:
     num_inducing = 7
@@ -28,3 +26,12 @@ def test_smoke_build_constant_input_dim_deep_gp(input_dim, num_layers):
     model_train.fit((X, Y), epochs=1)
     model_test = dgp.as_prediction_model()
     _ = model_test(X)
+
+
+@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.int32])
+def test_build_constant_input_dim_deep_gp_raises_on_incorrect_dtype(dtype):
+    config = make_dataclass_from_class(Config, DemoConfig)
+    X = np.random.randn(13, 2).astype(dtype)
+
+    with pytest.raises(ValueError):
+        build_constant_input_dim_deep_gp(X, 2, config)

tests/gpflux/layers/test_latent_variable_layer.py

@@ -26,8 +26,6 @@
 from gpflux.encoders import DirectlyParameterizedNormalDiag
 from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer
 
-tf.keras.backend.set_floatx("float64")
-
 ############
 # Utilities
 ############

tests/gpflux/models/test_bayesian_model.py

@@ -24,8 +24,6 @@
 from gpflux.layers import LatentVariableLayer, LikelihoodLayer
 from tests.integration.test_latent_variable_integration import build_gp_layers  # noqa: F401
 
-tf.keras.backend.set_floatx("float64")
-
 MAXITER = int(80e3)
 PLOTTER_INTERVAL = 60
 

tests/gpflux/models/test_deep_gp.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 import numpy as np
+import pytest
 import tensorflow as tf
 import tqdm
 
@@ -81,15 +82,15 @@ def step():
                 plotter()
 
 
-def setup_dataset(input_dim: int, num_data: int):
+def setup_dataset(input_dim: int, num_data: int, dtype: np.dtype = np.float64):
     lim = [0, 100]
     kernel = RBF(lengthscales=20)
     sigma = 0.01
     X = np.random.random(size=(num_data, input_dim)) * lim[1]
     cov = kernel.K(X) + np.eye(num_data) * sigma ** 2
     Y = np.random.multivariate_normal(np.zeros(num_data), cov)[:, None]
     Y = np.clip(Y, -0.5, 0.5)
-    return X, Y
+    return X.astype(dtype), Y.astype(dtype)
 
 
 def get_live_plotter(train_data, model):
@@ -133,7 +134,6 @@ def plotter(*args, **kwargs):
 
 
 def run_demo(maxiter=int(80e3), plotter_interval=60):
-    tf.keras.backend.set_floatx("float64")
     input_dim = 2
     num_data = 1000
     data = setup_dataset(input_dim, num_data)
@@ -155,6 +155,23 @@ def test_smoke():
     run_demo(maxiter=2, plotter_interval=1)
 
 
+@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.int32])
+def test_deep_gp_raises_on_incorrect_dtype(dtype):
+    input_dim = 2
+    num_data = 1000
+    X, Y = setup_dataset(input_dim, num_data, dtype)
+    model = build_deep_gp(input_dim, num_data)
+
+    with pytest.raises(ValueError):
+        model.predict_f(X)
+
+    with pytest.raises(ValueError):
+        model.call(X)
+
+    with pytest.raises(ValueError):
+        model.call(X, Y)
+
+
 if __name__ == "__main__":
     run_demo()
     input()

tests/gpflux/test_losses.py

@@ -7,8 +7,6 @@
 from gpflux.layers import LikelihoodLayer
 from gpflux.losses import LikelihoodLoss
 
-tf.keras.backend.set_floatx("float64")
-
 
 def test_likelihood_layer_and_likelihood_loss_give_equal_results():
     np.random.seed(123)

tests/integration/test_compilation.py

@@ -27,9 +27,6 @@
 from gpflux.losses import LikelihoodLoss
 from gpflux.models import DeepGP
 
-tf.keras.backend.set_floatx("float64")
-
-
 #########################################
 # Helpers
 #########################################

tests/integration/test_latent_variable_integration.py

@@ -30,8 +30,6 @@
 from gpflux.layers import GPLayer, LatentVariableLayer, LikelihoodLayer
 from gpflux.models import DeepGP
 
-tf.keras.backend.set_floatx("float64")
-
 ############
 # Utilities
 ############

tests/integration/test_svgp_equivalence.py

@@ -28,8 +28,6 @@
 
 import gpflux
 
-tf.keras.backend.set_floatx("float64")
-
 
 class LogPrior_ELBO_SVGP(gpflow.models.SVGP):
     """
@@ -264,15 +262,15 @@ def optimization_step():
 
 
 @pytest.mark.parametrize(
-    "svgp_fitter, sldgp_fitter",
+    "svgp_fitter, sldgp_fitter, tol_kw",
     [
-        (fit_adam, fit_adam),
-        (fit_adam, keras_fit_adam),
-        (fit_natgrad, fit_natgrad),
-        (fit_natgrad, keras_fit_natgrad),
+        (fit_adam, fit_adam, {}),
+        (fit_adam, keras_fit_adam, {}),
+        (fit_natgrad, fit_natgrad, {}),
+        (fit_natgrad, keras_fit_natgrad, dict(atol=1e-7)),
     ],
 )
-def test_svgp_equivalence_with_sldgp(svgp_fitter, sldgp_fitter, maxiter=20):
+def test_svgp_equivalence_with_sldgp(svgp_fitter, sldgp_fitter, tol_kw, maxiter=20):
     data = load_data()
 
     svgp = create_gpflow_svgp(*make_kernel_likelihood_iv())
@@ -281,14 +279,14 @@ def test_svgp_equivalence_with_sldgp(svgp_fitter, sldgp_fitter, maxiter=20):
     sldgp = create_gpflux_sldgp(*make_kernel_likelihood_iv(), get_num_data(data))
     sldgp_fitter(sldgp, data, maxiter=maxiter)
 
-    assert_equivalence(svgp, sldgp, data)
+    assert_equivalence(svgp, sldgp, data, **tol_kw)
 
 
 @pytest.mark.parametrize(
     "svgp_fitter, keras_fitter, tol_kw",
     [
         (fit_adam, _keras_fit_adam, {}),
-        (fit_natgrad, _keras_fit_natgrad, dict(atol=1e-8)),
+        (fit_natgrad, _keras_fit_natgrad, dict(atol=1e-6)),
     ],
 )
 def test_svgp_equivalence_with_keras_sequential(svgp_fitter, keras_fitter, tol_kw, maxiter=10):