typing: Use Self from typing_extensions

botorch/models/approximate_gp.py

@@ -32,7 +32,7 @@
 import copy
 import warnings
 
-from typing import Optional, TypeVar, Union
+from typing import Optional, Union
 
 import torch
 from botorch.exceptions.warnings import UserInputWarning
@@ -68,9 +68,9 @@
 )
 from torch import Tensor
 from torch.nn import Module
+from typing_extensions import Self
 
 
-TApproxModel = TypeVar("TApproxModel", bound="ApproximateGPyTorchModel")
 TRANSFORM_WARNING = (
     "Using an {ttype} transform with `SingleTaskVariationalGP`. If this "
     "model is trained in minibatches, a {ttype} transform with learnable "
@@ -132,11 +132,11 @@ def __init__(
     def num_outputs(self):
         return self._desired_num_outputs
 
-    def eval(self: TApproxModel) -> TApproxModel:
+    def eval(self) -> Self:
         r"""Puts the model in `eval` mode."""
         return Module.eval(self)
 
-    def train(self: TApproxModel, mode: bool = True) -> TApproxModel:
+    def train(self, mode: bool = True) -> Self:
         r"""Put the model in `train` mode.
 
         Args:

botorch/models/model.py

@@ -16,7 +16,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Mapping
-from typing import Any, Callable, Optional, TYPE_CHECKING, TypeVar, Union
+from typing import Any, Callable, Optional, TYPE_CHECKING, Union
 
 import numpy as np
 import torch
@@ -37,12 +37,11 @@
 from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood
 from torch import Tensor
 from torch.nn import Module, ModuleDict, ModuleList
+from typing_extensions import Self
 
 if TYPE_CHECKING:
     from botorch.acquisition.objective import PosteriorTransform  # pragma: no cover
 
-TFantasizeMixin = TypeVar("TFantasizeMixin", bound="FantasizeMixin")
-
 
 class Model(Module, ABC):
     r"""Abstract base class for BoTorch models.
@@ -289,11 +288,7 @@ def __init__(self, args):
     """
 
     @abstractmethod
-    def condition_on_observations(
-        self: TFantasizeMixin,
-        X: Tensor,
-        Y: Tensor,
-    ) -> TFantasizeMixin:
+    def condition_on_observations(self, X: Tensor, Y: Tensor) -> Self:
         """
         Classes that inherit from `FantasizeMixin` must implement
         a `condition_on_observations` method.
@@ -322,16 +317,13 @@ def transform_inputs(
         a `transform_inputs` method.
         """
 
-    # When Python 3.11 arrives we can start annotating return types like
-    # this as
-    # 'Self', but at this point the verbose 'T...' syntax is needed.
     def fantasize(
-        self: TFantasizeMixin,
+        self,
         X: Tensor,
         sampler: MCSampler,
         observation_noise: Optional[Tensor] = None,
         **kwargs: Any,
-    ) -> TFantasizeMixin:
+    ) -> Self:
         r"""Construct a fantasy model.
 
         Constructs a fantasy model in the following fashion:

requirements.txt

@@ -3,5 +3,6 @@ scipy
 mpmath>=0.19,<=1.3
 torch>=2.0.1
 pyro-ppl>=1.8.4
+typing_extensions
 gpytorch==1.13
 linear_operator==0.5.3