Use correct dtype for eps in IndependentReparametrizationSampler (#824)

tests/integration/models/multifidelity/test_multifidelity_models.py

@@ -5,6 +5,7 @@
 from tensorflow.keras.metrics import mean_squared_error
 
 import trieste
+from tests.util.misc import random_seed
 from trieste.data import (
     Dataset,
     add_fidelity_column,
@@ -144,6 +145,7 @@ def test_multifidelity_nonlinear_autoregressive_results_better_than_linear() ->
     assert mses[0] < mses[1]
 
 
+@random_seed
 def test_multifidelity_autoregressive_gets_expected_rhos() -> None:
     input_dim = 1
     lb = np.zeros(input_dim)

tests/unit/models/gpflow/test_sampler.py

@@ -252,6 +252,19 @@ def test_independent_reparametrization_sampler_samples_are_distinct_for_new_inst
         npt.assert_array_less(1e-9, tf.abs(sampler2.sample(xs) - sampler1.sample(xs)))
 
 
+@pytest.mark.parametrize("qmc", [True, False])
+@pytest.mark.parametrize("qmc_skip", [True, False])
+@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
+def test_independent_reparametrization_sampler_dtype(
+    qmc: bool, qmc_skip: bool, dtype: tf.DType
+) -> None:
+    model = QuadraticMeanAndRBFKernel()
+    sampler = IndependentReparametrizationSampler(2, model, qmc=qmc, qmc_skip=qmc_skip)
+    xs = tf.random.uniform([5, 1, 2], minval=-10.0, maxval=10.0, dtype=dtype)
+    samples = sampler.sample(xs)
+    assert samples.dtype is dtype
+
+
 @pytest.mark.parametrize("qmc", [True, False])
 @pytest.mark.parametrize("qmc_skip", [True, False])
 def test_independent_reparametrization_sampler_reset_sampler(qmc: bool, qmc_skip: bool) -> None:
@@ -885,12 +898,19 @@ def test_rff_and_decoupled_trajectory_give_similar_results(
     )  # variance across samples should (very) roughly agree for different samplers
 
 
-@pytest.mark.parametrize("n_sample_dim", [2, 5])
+@pytest.mark.parametrize("n_sample_dim", [0, 2, 5])
 @pytest.mark.parametrize("skip", [0, 10_000])
-def test_qmc_samples_return_standard_normal_samples(n_sample_dim: int, skip: int) -> None:
+@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
+def test_qmc_samples_return_standard_normal_samples(
+    n_sample_dim: int, skip: int, dtype: tf.DType
+) -> None:
     n_samples = 10_000
 
-    qmc_samples = qmc_normal_samples(num_samples=n_samples, n_sample_dim=n_sample_dim, skip=skip)
+    qmc_samples = qmc_normal_samples(
+        num_samples=n_samples, n_sample_dim=n_sample_dim, skip=skip, dtype=dtype
+    )
+    assert qmc_samples.dtype is dtype
+    assert qmc_samples.shape == (n_samples, n_sample_dim)
 
     # should be multivariate normal with zero correlation
     for i in range(n_sample_dim):

trieste/models/gpflow/sampler.py

@@ -49,26 +49,29 @@
 
 
 def qmc_normal_samples(
-    num_samples: _IntTensorType, n_sample_dim: _IntTensorType, skip: _IntTensorType = 0
+    num_samples: _IntTensorType,
+    n_sample_dim: _IntTensorType,
+    skip: _IntTensorType = 0,
+    dtype: tf.DType = tf.float64,
 ) -> tf.Tensor:
     """
     Generates `num_samples` sobol samples, skipping the first `skip`, where each
     sample has dimension `n_sample_dim`.
     """
 
     if num_samples == 0 or n_sample_dim == 0:
-        return tf.zeros(shape=(num_samples, n_sample_dim), dtype=tf.float64)
+        return tf.zeros(shape=(num_samples, n_sample_dim), dtype=dtype)
 
     sobol_samples = tf.math.sobol_sample(
         dim=n_sample_dim,
         num_results=num_samples,
-        dtype=tf.float64,
+        dtype=dtype,
         skip=skip,
     )
 
     dist = tfp.distributions.Normal(
-        loc=tf.constant(0.0, dtype=tf.float64),
-        scale=tf.constant(1.0, dtype=tf.float64),
+        loc=tf.constant(0.0, dtype=dtype),
+        scale=tf.constant(1.0, dtype=dtype),
     )
     normal_samples = dist.quantile(sobol_samples)
     return normal_samples
@@ -138,10 +141,12 @@ def sample_eps() -> tf.Tensor:
                     IndependentReparametrizationSampler.skip.assign(skip + self._sample_size)
                 else:
                     skip = tf.constant(0)
-                normal_samples = qmc_normal_samples(self._sample_size, mean.shape[-1], skip)
+                normal_samples = qmc_normal_samples(
+                    self._sample_size, mean.shape[-1], skip, dtype=var.dtype
+                )
             else:
                 normal_samples = tf.random.normal(
-                    [self._sample_size, tf.shape(mean)[-1]], dtype=tf.float64
+                    [self._sample_size, tf.shape(mean)[-1]], dtype=var.dtype
                 )
             return normal_samples  # [S, L]
 
@@ -154,7 +159,7 @@ def sample_eps() -> tf.Tensor:
             lambda: self._eps.assign(sample_eps()),
         )
 
-        return mean + tf.sqrt(var) * tf.cast(self._eps[:, None, :], var.dtype)  # [..., S, 1, L]
+        return mean + tf.sqrt(var) * self._eps[:, None, :]  # [..., S, 1, L]
 
 
 class BatchReparametrizationSampler(ReparametrizationSampler[SupportsPredictJoint]):