Remove superfluous encodings

secondmind-labs · Jul 31, 2024 · c071112 · c071112
1 parent 25d20ca
commit c071112
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 12 deletions.
diff --git a/trieste/models/gpflow/models.py b/trieste/models/gpflow/models.py
@@ -50,7 +50,6 @@
     SupportsPredictY,
     TrainableProbabilisticModel,
     TrajectorySampler,
-    encode_query_points,
 )
 from ..optimizer import BatchOptimizer, Optimizer, OptimizeResult
 from .inducing_point_selectors import InducingPointSelector
@@ -165,7 +164,6 @@ def _ensure_variable_model_data(self) -> None:
         )
 
     @inherit_check_shapes
-    @encode_query_points
     def predict_y(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
         f_mean, f_var = self.predict(query_points)
         return self.model.likelihood.predict_mean_and_var(query_points, f_mean, f_var)
@@ -355,7 +353,6 @@ def get_internal_data(self) -> Dataset:
         """
         return Dataset(self.model.data[0], self.model.data[1])
 
-    @encode_query_points
     def conditional_predict_f(
         self, query_points: TensorType, additional_data: Dataset
     ) -> tuple[TensorType, TensorType]:
@@ -421,7 +418,6 @@ def conditional_predict_f(
 
         return mean_qp_new, var_qp_new
 
-    @encode_query_points
     def conditional_predict_joint(
         self, query_points: TensorType, additional_data: Dataset
     ) -> tuple[TensorType, TensorType]:
@@ -492,7 +488,6 @@ def conditional_predict_joint(
 
         return mean_qp_new, cov_qp_new
 
-    @encode_query_points
     def conditional_predict_f_sample(
         self, query_points: TensorType, additional_data: Dataset, num_samples: int
     ) -> TensorType:
@@ -514,7 +509,6 @@ def conditional_predict_f_sample(
         )  # [..., (S), P, N]
         return tf.linalg.adjoint(samples)  # [..., (S), N, L]
 
-    @encode_query_points
     def conditional_predict_y(
         self, query_points: TensorType, additional_data: Dataset
     ) -> tuple[TensorType, TensorType]:
@@ -623,7 +617,6 @@ def inducing_point_selector(
         return self._inducing_point_selector
 
     @inherit_check_shapes
-    @encode_query_points
     def predict_y(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
         f_mean, f_var = self.predict(query_points)
         return self.model.likelihood.predict_mean_and_var(query_points, f_mean, f_var)
@@ -950,7 +943,6 @@ def inducing_point_selector(self) -> Optional[InducingPointSelector[SparseVariat
         return self._inducing_point_selector
 
     @inherit_check_shapes
-    @encode_query_points
     def predict_y(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
         f_mean, f_var = self.predict(query_points)
         return self.model.likelihood.predict_mean_and_var(query_points, f_mean, f_var)
@@ -1267,7 +1259,6 @@ def model(self) -> VGP:
         return self._model
 
     @inherit_check_shapes
-    @encode_query_points
     def predict_y(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
         f_mean, f_var = self.predict(query_points)
         return self.model.likelihood.predict_mean_and_var(query_points, f_mean, f_var)

diff --git a/trieste/models/interfaces.py b/trieste/models/interfaces.py
@@ -765,8 +765,9 @@ def encoder(self) -> EncoderFunction | None:
 
 
 def encode_query_points(f: C) -> C:
-    """Decorator for encoding query points, applicable to HasEncoder model methods such as predict
-    whose first argument is query_points.
+    """
+    Decorator for automatically encoding query points, applicable to HasEncoder model methods
+    such as predict whose first argument is query_points.
     """
 
     @functools.wraps(f)

diff --git a/trieste/models/keras/models.py b/trieste/models/keras/models.py
@@ -332,7 +332,6 @@ def predict_ensemble(self, query_points: TensorType) -> tuple[TensorType, Tensor
         return predicted_means, predicted_vars
 
     @inherit_check_shapes
-    @encode_query_points
     def sample(self, query_points: TensorType, num_samples: int) -> TensorType:
         """
         Return ``num_samples`` samples at ``query_points``. We use the mixture approximation in