Really leave them alone

secondmind-labs · Jan 6, 2025 · bfccfd8 · bfccfd8
1 parent 51adc46
commit bfccfd8
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 18 deletions.
diff --git a/tests/unit/models/gpflow/test_sampler.py b/tests/unit/models/gpflow/test_sampler.py
@@ -471,20 +471,6 @@ def test_batch_reparametrization_sampler_reset_sampler(qmc: bool, qmc_skip: bool
         npt.assert_array_less(1e-9, tf.abs(samples2 - samples1))
 
 
-@pytest.mark.parametrize("qmc", [True, False])
-@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
-def test_batch_reparametrization_sampler_sample_ensures_positive_variance(
-    qmc: bool, dtype: tf.DType
-) -> None:
-    model = QuadraticMeanAndRBFKernel(kernel_amplitude=tf.constant(0, dtype=dtype))
-    sampler = BatchReparametrizationSampler(100, model, qmc=qmc)
-    x = tf.constant([[1.0]], dtype=dtype)
-    variance = tf.math.reduce_variance(sampler.sample(x))  # default jitter
-    assert variance > 1e-7
-    variance = tf.math.reduce_variance(sampler.sample(x, jitter=0.0))  # explicit jitter
-    assert variance > (1e-7 if dtype is tf.float32 else 1e-17)
-
-
 @pytest.mark.parametrize("num_features", [0, -2])
 def test_rff_trajectory_sampler_raises_for_invalid_number_of_features(
     num_features: int,

diff --git a/trieste/models/gpflow/sampler.py b/trieste/models/gpflow/sampler.py
@@ -277,8 +277,7 @@ def sample_eps() -> tf.Tensor:
             )
 
         identity = tf.eye(batch_size, dtype=cov.dtype)  # [B, B]
-        cov = ensure_positive(cov + jitter * identity)
-        cov_cholesky = tf.linalg.cholesky(cov)  # [..., L, B, B]
+        cov_cholesky = tf.linalg.cholesky(cov + jitter * identity)  # [..., L, B, B]
 
         variance_contribution = cov_cholesky @ self._eps  # [..., L, B, S]
 

diff --git a/trieste/models/gpflux/sampler.py b/trieste/models/gpflux/sampler.py
@@ -27,7 +27,6 @@
 
 from ...types import TensorType
 from ...utils import DEFAULTS, flatten_leading_dims
-from ...utils.misc import ensure_positive
 from ..interfaces import (
     ReparametrizationSampler,
     TrajectoryFunction,
@@ -109,7 +108,7 @@ def sample(self, at: TensorType, *, jitter: float = DEFAULTS.JITTER) -> TensorTy
                 continue
 
             mean, var = layer.predict(samples, full_cov=False, full_output_cov=False)
-            var = ensure_positive(var + jitter)
+            var = var + jitter
 
             if not self._initialized:
                 self._eps_list[i].assign(