Expand Self usage for better fluent APIs (#3194)

Summary:
Pull Request resolved: #3194

Replace explicit class return types with `Self` for methods that return
`self`, enabling better type inference for subclasses.

Files updated:
- botorch/models/relevance_pursuit.py
  - RelevancePursuitMixin: to_sparse, to_dense, expand_support, contract_support,
    full_support, remove_support now return Self
- botorch/models/robust_relevance_pursuit_model.py
  - RobustRelevancePursuitMixin.load_standard_model now returns Self
- botorch/models/model.py
  - Model.eval, Model.train now return Self

Reviewed By: hvarfner

Differential Revision: D91641969

fbshipit-source-id: 6bac972ad533ebb0c23a86e9de36cb6cfb99eaf9

Author: Balandat

botorch/models/model.py

@@ -243,12 +243,12 @@ def _revert_to_original_inputs(self) -> None:
             self.set_train_data(self._original_train_inputs, strict=False)
             self._has_transformed_inputs = False
 
-    def eval(self) -> Model:
+    def eval(self) -> Self:
         r"""Puts the model in ``eval`` mode and sets the transformed inputs."""
         self._set_transformed_inputs()
         return super().eval()
 
-    def train(self, mode: bool = True) -> Model:
+    def train(self, mode: bool = True) -> Self:
         r"""Put the model in ``train`` mode. Reverts to the original inputs if
         in ``train`` mode (``mode=True``) or sets transformed inputs if in
         ``eval`` mode (``mode=False``).

botorch/models/relevance_pursuit.py

@@ -22,7 +22,7 @@
 from collections.abc import Callable, Sequence
 from copy import copy, deepcopy
 from functools import partial
-from typing import Any, cast
+from typing import Any, cast, Self
 from warnings import warn
 
 import torch
@@ -117,7 +117,7 @@ def inactive_indices(self) -> Tensor:
         device = self.sparse_parameter.device
         return torch.arange(self.dim, device=device)[~self.is_active]
 
-    def to_sparse(self) -> RelevancePursuitMixin:
+    def to_sparse(self) -> Self:
         """Converts the sparse parameter to its sparse (< dim) representation.
 
         Returns:
@@ -130,7 +130,7 @@ def to_sparse(self) -> RelevancePursuitMixin:
             self._is_sparse = True
         return self
 
-    def to_dense(self) -> RelevancePursuitMixin:
+    def to_dense(self) -> Self:
         """Converts the sparse parameter to its dense, length-``dim`` representation.
 
         Returns:
@@ -157,7 +157,7 @@ def to_dense(self) -> RelevancePursuitMixin:
             self._is_sparse = False
         return self
 
-    def expand_support(self, indices: list[int]) -> RelevancePursuitMixin:
+    def expand_support(self, indices: list[int]) -> Self:
         """Expands the support by a number of indices.
 
         Args:
@@ -186,7 +186,7 @@ def expand_support(self, indices: list[int]) -> RelevancePursuitMixin:
             )
         return self
 
-    def contract_support(self, indices: list[int]) -> RelevancePursuitMixin:
+    def contract_support(self, indices: list[int]) -> Self:
         """Contracts the support by a number of indices.
 
         Args:
@@ -216,7 +216,7 @@ def contract_support(self, indices: list[int]) -> RelevancePursuitMixin:
         return self
 
     # support initialization helpers
-    def full_support(self) -> RelevancePursuitMixin:
+    def full_support(self) -> Self:
         """Initializes the RelevancePursuitMixin with a full, size-``dim`` support.
 
         Returns:
@@ -226,7 +226,7 @@ def full_support(self) -> RelevancePursuitMixin:
         self.to_dense()  # no reason to be sparse with full support
         return self
 
-    def remove_support(self) -> RelevancePursuitMixin:
+    def remove_support(self) -> Self:
         """Initializes the RelevancePursuitMixin with an empty, size-zero support.
 
         Returns:

botorch/models/robust_relevance_pursuit_model.py

@@ -34,7 +34,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Mapping, Sequence
-from typing import Any
+from typing import Any, Self
 
 import torch
 from botorch.exceptions.errors import UnsupportedError
@@ -140,7 +140,7 @@ def to_standard_model(self) -> Model:
             A standard model.
         """
 
-    def load_standard_model(self, standard_model: Model) -> RobustRelevancePursuitMixin:
+    def load_standard_model(self, standard_model: Model) -> Self:
         """Loads the state dict of a model into the ``RobustRelevancePursuitMixin``.
 
         Args: