minimum not maximum

secondmind-labs · Jan 3, 2025 · e284c87 · e284c87
1 parent 8ef3e46
commit e284c87
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/trieste/models/gpflow/sampler.py b/trieste/models/gpflow/sampler.py
@@ -133,7 +133,7 @@ def sample(self, at: TensorType, *, jitter: float = DEFAULTS.JITTER) -> TensorTy
         tf.debugging.assert_greater_equal(jitter, 0.0)
 
         mean, var = self._model.predict(at[..., None, :, :])  # [..., 1, 1, L], [..., 1, 1, L]
-        var = var + tf.math.maximum(var, jitter)
+        var = var + tf.math.minimum(var, jitter)
 
         def sample_eps() -> tf.Tensor:
             self._initialized.assign(True)
@@ -276,7 +276,7 @@ def sample_eps() -> tf.Tensor:
             )
 
         identity = tf.eye(batch_size, dtype=cov.dtype)  # [B, B]
-        cov = cov + tf.math.maximum(cov, jitter) * identity
+        cov = cov + tf.math.minimum(cov, jitter) * identity
         cov_cholesky = tf.linalg.cholesky(cov)  # [..., 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
@@ -108,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 = var + tf.math.maximum(var, jitter)
+            var = var + tf.math.minimum(var, jitter)
 
             if not self._initialized:
                 self._eps_list[i].assign(