secondmind-labs · hstojic · Sep 2, 2022 · Aug 17, 2022 · Aug 17, 2022 · Aug 31, 2022
diff --git a/benchmarking/main.py b/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")

diff --git a/docs/notebooks/deep_cde.ipynb b/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": {}

diff --git a/docs/notebooks/efficient_sampling.py b/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]
 """

diff --git a/docs/notebooks/gpflux_features.py b/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")
 
 

diff --git a/docs/notebooks/gpflux_with_keras_layers.py b/docs/notebooks/gpflux_with_keras_layers.py
@@ -30,7 +30,6 @@
 
 from gpflow.config import default_float
 
-tf.keras.backend.set_floatx("float64")
 
 # %% [markdown]
 """

diff --git a/docs/notebooks/intro.py b/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]

diff --git a/docs/notebooks/keras_integration.py b/docs/notebooks/keras_integration.py
@@ -28,8 +28,6 @@
 
 import matplotlib.pyplot as plt
 
-# %%
-tf.keras.backend.set_floatx("float64")
 
 # %%
 # %matplotlib inline

diff --git a/gpflux/layers/latent_variable_layer.py b/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 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,7 @@ 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(

diff --git a/gpflux/models/deep_gp.py b/gpflux/models/deep_gp.py
@@ -96,8 +96,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=tf.float64, name="inputs")
+        self.targets = tf.keras.Input((target_dim,), dtype=tf.float64, name="targets")
         self.f_layers = f_layers
         if isinstance(likelihood, gpflow.likelihoods.Likelihood):
             self.likelihood_layer = LikelihoodLayer(likelihood)

diff --git a/tests/gpflux/architectures/test_constant_input_dim_deep_gp.py b/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

diff --git a/tests/gpflux/layers/test_latent_variable_layer.py b/tests/gpflux/layers/test_latent_variable_layer.py
@@ -26,7 +26,6 @@
 from gpflux.encoders import DirectlyParameterizedNormalDiag
 from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer
 
-tf.keras.backend.set_floatx("float64")
 
 ############
 # Utilities

diff --git a/tests/gpflux/models/test_bayesian_model.py b/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
 

diff --git a/tests/gpflux/models/test_deep_gp.py b/tests/gpflux/models/test_deep_gp.py
@@ -133,7 +133,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)

diff --git a/tests/gpflux/test_losses.py b/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)

diff --git a/tests/integration/test_compilation.py b/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
 #########################################

diff --git a/tests/integration/test_latent_variable_integration.py b/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
 ############

diff --git a/tests/integration/test_svgp_equivalence.py b/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):