Make X, Y, Yvar into properties in datasets (#2004)

Summary:
X-link: facebook/Ax#1817

Pull Request resolved: #2004

Since the primary use case of the datasets is to house a couple tensors, this simplifies the process by making dataset.X/Y/Yvar into a tensor rather than a callable.

This also eliminates the need to make tensors into containers, eliminating a step in the process.

Reviewed By: Balandat

Differential Revision: D48926544

fbshipit-source-id: ca31fd0946c449b501bb5a9e5d3f8f5e38bab715

Author: saitcakmak

botorch/acquisition/input_constructors.py

@@ -141,7 +141,10 @@ def _field_is_shared(
         obj = getattr(dataset, fieldname)
         if base is None:
             base = obj
-        elif base != obj:
+        elif isinstance(base, Tensor):
+            if not torch.equal(base, obj):
+                return False
+        elif base != obj:  # pragma: no cover
             return False
 
     return True
@@ -414,7 +417,7 @@ def construct_inputs_noisy_ei(
     X = _get_dataset_field(training_data, "X", first_only=True, assert_shared=True)
     return {
         "model": model,
-        "X_observed": X(),
+        "X_observed": X,
         "num_fantasies": num_fantasies,
         "maximize": maximize,
     }
@@ -624,7 +627,6 @@ def construct_inputs_qNEI(
         X_baseline = _get_dataset_field(
             training_data,
             fieldname="X",
-            transform=lambda field: field(),
             assert_shared=True,
             first_only=True,
         )
@@ -845,7 +847,6 @@ def construct_inputs_EHVI(
     X = _get_dataset_field(
         training_data,
         fieldname="X",
-        transform=lambda field: field(),
         first_only=True,
         assert_shared=True,
     )
@@ -908,7 +909,6 @@ def construct_inputs_qEHVI(
     X = _get_dataset_field(
         training_data,
         fieldname="X",
-        transform=lambda field: field(),
         first_only=True,
         assert_shared=True,
     )
@@ -974,7 +974,6 @@ def construct_inputs_qNEHVI(
         X_baseline = _get_dataset_field(
             training_data,
             fieldname="X",
-            transform=lambda field: field(),
             first_only=True,
             assert_shared=True,
         )
@@ -1245,7 +1244,6 @@ def get_best_f_analytic(
     Y = _get_dataset_field(
         training_data,
         fieldname="Y",
-        transform=lambda field: field(),
         join_rule=lambda field_tensors: torch.cat(field_tensors, dim=-1),
     )
 
@@ -1274,15 +1272,13 @@ def get_best_f_mc(
     X_baseline = _get_dataset_field(
         training_data,
         fieldname="X",
-        transform=lambda field: field(),
         assert_shared=True,
         first_only=True,
     )
 
     Y = _get_dataset_field(
         training_data,
         fieldname="Y",
-        transform=lambda field: field(),
         join_rule=lambda field_tensors: torch.cat(field_tensors, dim=-1),
     )  # batch_shape x n x d
 

botorch/models/utils/parse_training_data.py

@@ -51,19 +51,21 @@ def parse_training_data(
 def _parse_model_supervised(
     consumer: Model, dataset: SupervisedDataset, **ignore: Any
 ) -> Dict[str, Tensor]:
-    parsed_data = {"train_X": dataset.X(), "train_Y": dataset.Y()}
+    parsed_data = {"train_X": dataset.X, "train_Y": dataset.Y}
     if dataset.Yvar is not None:
-        parsed_data["train_Yvar"] = dataset.Yvar()
+        parsed_data["train_Yvar"] = dataset.Yvar
     return parsed_data
 
 
 @dispatcher.register(PairwiseGP, RankingDataset)
 def _parse_pairwiseGP_ranking(
     consumer: PairwiseGP, dataset: RankingDataset, **ignore: Any
 ) -> Dict[str, Tensor]:
-    datapoints = dataset.X.values
-    comparisons = dataset.X.indices
-    comp_order = dataset.Y()
+    # TODO: [T163045056] Not sure what the point of the special container is if we have
+    # to further process it here. We should move this logic into RankingDataset.
+    datapoints = dataset._X.values
+    comparisons = dataset._X.indices
+    comp_order = dataset.Y
     comparisons = torch.gather(input=comparisons, dim=-1, index=comp_order)
 
     return {

botorch/utils/datasets.py

@@ -12,11 +12,11 @@
 from itertools import count, repeat
 from typing import Any, Dict, Hashable, Iterable, Optional, TypeVar, Union
 
-from botorch.utils.containers import BotorchContainer, DenseContainer, SliceContainer
+import torch
+from botorch.utils.containers import BotorchContainer, SliceContainer
 from torch import long, ones, Tensor
 
 T = TypeVar("T")
-ContainerLike = Union[BotorchContainer, Tensor]
 MaybeIterable = Union[T, Iterable[T]]
 
 
@@ -25,9 +25,6 @@ class SupervisedDataset:
     and an optional `Yvar` that stipulates observations variances so
     that `Y[i] ~ N(f(X[i]), Yvar[i])`.
 
-    This class object's `__init__` method converts Tensors `src` to
-    DenseContainers under the assumption that `event_shape=src.shape[-1:]`.
-
     Example:
 
     .. code-block:: python
@@ -42,15 +39,11 @@ class SupervisedDataset:
         assert A == B
     """
 
-    X: BotorchContainer
-    Y: BotorchContainer
-    Yvar: Optional[BotorchContainer]
-
     def __init__(
         self,
-        X: ContainerLike,
-        Y: ContainerLike,
-        Yvar: Optional[ContainerLike] = None,
+        X: Union[BotorchContainer, Tensor],
+        Y: Union[BotorchContainer, Tensor],
+        Yvar: Union[BotorchContainer, Tensor, None] = None,
         validate_init: bool = True,
     ) -> None:
         r"""Constructs a `SupervisedDataset`.
@@ -62,17 +55,41 @@ def __init__(
                 the observation noise.
             validate_init: If `True`, validates the input shapes.
         """
-        self.X = _containerize(X)
-        self.Y = _containerize(Y)
-        self.Yvar = None if Yvar is None else _containerize(Yvar)
+        self._X = X
+        self._Y = Y
+        self._Yvar = Yvar
         if validate_init:
             self._validate()
 
+    @property
+    def X(self) -> Tensor:
+        if isinstance(self._X, Tensor):
+            return self._X
+        return self._X()
+
+    @property
+    def Y(self) -> Tensor:
+        if isinstance(self._Y, Tensor):
+            return self._Y
+        return self._Y()
+
+    @property
+    def Yvar(self) -> Optional[Tensor]:
+        if self._Yvar is None or isinstance(self._Yvar, Tensor):
+            return self._Yvar
+        return self._Yvar()
+
     def _validate(self) -> None:
         shape_X = self.X.shape
-        shape_X = shape_X[: len(shape_X) - len(self.X.event_shape)]
+        if isinstance(self._X, BotorchContainer):
+            shape_X = shape_X[: len(shape_X) - len(self._X.event_shape)]
+        else:
+            shape_X = shape_X[:-1]
         shape_Y = self.Y.shape
-        shape_Y = shape_Y[: len(shape_Y) - len(self.Y.event_shape)]
+        if isinstance(self._Y, BotorchContainer):
+            shape_Y = shape_Y[: len(shape_Y) - len(self._Y.event_shape)]
+        else:
+            shape_Y = shape_Y[:-1]
         if shape_X != shape_Y:
             raise ValueError("Batch dimensions of `X` and `Y` are incompatible.")
         if self.Yvar is not None and self.Yvar.shape != self.Y.shape:
@@ -81,9 +98,9 @@ def _validate(self) -> None:
     @classmethod
     def dict_from_iter(
         cls,
-        X: MaybeIterable[ContainerLike],
-        Y: MaybeIterable[ContainerLike],
-        Yvar: Optional[MaybeIterable[ContainerLike]] = None,
+        X: MaybeIterable[Union[BotorchContainer, Tensor]],
+        Y: MaybeIterable[Union[BotorchContainer, Tensor]],
+        Yvar: Optional[MaybeIterable[Union[BotorchContainer, Tensor]]] = None,
         *,
         keys: Optional[Iterable[Hashable]] = None,
     ) -> Dict[Hashable, SupervisedDataset]:
@@ -106,9 +123,13 @@ def dict_from_iter(
     def __eq__(self, other: Any) -> bool:
         return (
             type(other) is type(self)
-            and self.X == other.X
-            and self.Y == other.Y
-            and self.Yvar == other.Yvar
+            and torch.equal(self.X, other.X)
+            and torch.equal(self.Y, other.Y)
+            and (
+                other.Yvar is None
+                if self.Yvar is None
+                else torch.equal(self.Yvar, other.Yvar)
+            )
         )
 
 
@@ -121,9 +142,9 @@ class FixedNoiseDataset(SupervisedDataset):
 
     def __init__(
         self,
-        X: ContainerLike,
-        Y: ContainerLike,
-        Yvar: ContainerLike,
+        X: Union[BotorchContainer, Tensor],
+        Y: Union[BotorchContainer, Tensor],
+        Yvar: Union[BotorchContainer, Tensor],
         validate_init: bool = True,
     ) -> None:
         r"""Initialize a `FixedNoiseDataset` -- deprecated!"""
@@ -159,11 +180,11 @@ class RankingDataset(SupervisedDataset):
         dataset = RankingDataset(X, Y)
     """
 
-    X: SliceContainer
-    Y: BotorchContainer
-
     def __init__(
-        self, X: SliceContainer, Y: ContainerLike, validate_init: bool = True
+        self,
+        X: SliceContainer,
+        Y: Union[BotorchContainer, Tensor],
+        validate_init: bool = True,
     ) -> None:
         r"""Construct a `RankingDataset`.
 
@@ -177,8 +198,8 @@ def __init__(
     def _validate(self) -> None:
         super()._validate()
 
-        Y = self.Y()
-        arity = self.X.indices.shape[-1]
+        Y = self.Y
+        arity = self._X.indices.shape[-1]
         if Y.min() < 0 or Y.max() >= arity:
             raise ValueError("Invalid ranking(s): out-of-bounds ranks detected.")
 
@@ -202,13 +223,3 @@ def _validate(self) -> None:
 
             # Same as: torch.where(y_diff == 0, y_incr + 1, 1)
             y_incr = y_incr - y_diff + 1
-
-
-def _containerize(value: ContainerLike) -> BotorchContainer:
-    r"""Converts Tensor-valued arguments to DenseContainer under the assumption
-    that said arguments house collections of feature vectors.
-    """
-    if isinstance(value, Tensor):
-        return DenseContainer(value, event_shape=value.shape[-1:])
-    else:
-        return value