Upgrade to use py3.10 features (#2712)

Author: Balandat

.pre-commit-config.yaml

@@ -11,6 +11,12 @@ repos:
         args: [--fix=lf]
     -   id: trailing-whitespace
     -   id: debug-statements
+-   repo: https://github.com/asottile/pyupgrade
+    rev: v3.19.1
+    hooks:
+    -   id: pyupgrade
+        args: [--py310-plus]
+        exclude: ^(examples/.*)|(docs/.*)
 -   repo: https://github.com/pycqa/flake8
     rev: 7.3.0
     hooks:

gpytorch/__init__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from typing import Optional, Tuple, Union
+from __future__ import annotations
 
 import linear_operator
 import torch
@@ -28,7 +28,7 @@
 from .mlls import ExactMarginalLogLikelihood
 from .module import Module
 
-Anysor = Union[LinearOperator, Tensor]
+Anysor = LinearOperator | Tensor
 
 
 def add_diagonal(input: Anysor, diag: Tensor) -> LinearOperator:
@@ -58,7 +58,7 @@ def add_jitter(input: Anysor, jitter_val: float = 1e-3) -> Anysor:
     return linear_operator.add_jitter(input=input, jitter_val=jitter_val)
 
 
-def diagonalization(input: Anysor, method: Optional[str] = None) -> Tuple[Tensor, Tensor]:
+def diagonalization(input: Anysor, method: str | None = None) -> tuple[Tensor, Tensor]:
     r"""
     Returns a (usually partial) diagonalization of a symmetric positive definite matrix (or batch of matrices).
     :math:`\mathbf A`.
@@ -74,7 +74,7 @@ def diagonalization(input: Anysor, method: Optional[str] = None) -> Tuple[Tensor
 
 
 def dsmm(
-    sparse_mat: Union[torch.sparse.HalfTensor, torch.sparse.FloatTensor, torch.sparse.DoubleTensor],
+    sparse_mat: torch.sparse.HalfTensor | torch.sparse.FloatTensor | torch.sparse.DoubleTensor,
     dense_mat: Tensor,
 ) -> Tensor:
     r"""
@@ -117,10 +117,10 @@ def inv_quad(input: Anysor, inv_quad_rhs: Tensor, reduce_inv_quad: bool = True)
 
 def inv_quad_logdet(
     input: Anysor,
-    inv_quad_rhs: Optional[Tensor] = None,
+    inv_quad_rhs: Tensor | None = None,
     logdet: bool = False,
     reduce_inv_quad: bool = True,
-) -> Tuple[Tensor, Tensor]:
+) -> tuple[Tensor, Tensor]:
     r"""
     Calls both :func:`inv_quad_logdet` and :func:`logdet` on a positive definite matrix (or batch) :math:`\mathbf A`.
     However, calling this method is far more efficient and stable than calling each method independently.
@@ -146,9 +146,9 @@ def inv_quad_logdet(
 def pivoted_cholesky(
     input: Anysor,
     rank: int,
-    error_tol: Optional[float] = None,
+    error_tol: float | None = None,
     return_pivots: bool = False,
-) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+) -> Tensor | tuple[Tensor, Tensor]:
     r"""
     Performs a partial pivoted Cholesky factorization of a positive definite matrix (or batch of matrices).
     :math:`\mathbf L \mathbf L^\top = \mathbf A`.
@@ -173,7 +173,7 @@ def pivoted_cholesky(
     return linear_operator.pivoted_cholesky(input=input, rank=rank, return_pivots=return_pivots)
 
 
-def root_decomposition(input: Anysor, method: Optional[str] = None) -> LinearOperator:
+def root_decomposition(input: Anysor, method: str | None = None) -> LinearOperator:
     r"""
     Returns a (usually low-rank) root decomposition linear operator of the
     positive definite matrix (or batch of matrices) :math:`\mathbf A`.
@@ -190,9 +190,9 @@ def root_decomposition(input: Anysor, method: Optional[str] = None) -> LinearOpe
 
 def root_inv_decomposition(
     input: Anysor,
-    initial_vectors: Optional[Tensor] = None,
-    test_vectors: Optional[Tensor] = None,
-    method: Optional[str] = None,
+    initial_vectors: Tensor | None = None,
+    test_vectors: Tensor | None = None,
+    method: str | None = None,
 ) -> LinearOperator:
     r"""
     Returns a (usually low-rank) inverse root decomposition linear operator
@@ -217,7 +217,7 @@ def root_inv_decomposition(
     )
 
 
-def solve(input: Anysor, rhs: Tensor, lhs: Optional[Tensor] = None) -> Tensor:
+def solve(input: Anysor, rhs: Tensor, lhs: Tensor | None = None) -> Tensor:
     r"""
     Given a positive definite matrix (or batch of matrices) :math:`\mathbf A`,
     computes a linear solve with right hand side :math:`\mathbf R`:
@@ -249,7 +249,7 @@ def solve(input: Anysor, rhs: Tensor, lhs: Optional[Tensor] = None) -> Tensor:
     return linear_operator.solve(input=input, rhs=rhs, lhs=lhs)
 
 
-def sqrt_inv_matmul(input: Anysor, rhs: Tensor, lhs: Optional[Tensor] = None) -> Tensor:
+def sqrt_inv_matmul(input: Anysor, rhs: Tensor, lhs: Tensor | None = None) -> Tensor:
     r"""
     Given a positive definite matrix (or batch of matrices) :math:`\mathbf A`
     and a right hand size :math:`\mathbf R`,

gpytorch/beta_features.py

@@ -1,18 +1,20 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import warnings
 
 from .settings import _feature_flag, _value_context
 
 
-class _moved_beta_feature(object):
+class _moved_beta_feature:
     def __init__(self, new_cls, orig_name=None):
         self.new_cls = new_cls
-        self.orig_name = orig_name if orig_name is not None else "gpytorch.settings.{}".format(new_cls.__name__)
+        self.orig_name = orig_name if orig_name is not None else f"gpytorch.settings.{new_cls.__name__}"
 
     def __call__(self, *args, **kwargs):
         warnings.warn(
-            "`{}` has moved to `gpytorch.settings.{}`.".format(self.orig_name, self.new_cls.__name__),
+            f"`{self.orig_name}` has moved to `gpytorch.settings.{self.new_cls.__name__}`.",
             DeprecationWarning,
         )
         return self.new_cls(*args, **kwargs)
@@ -55,7 +57,5 @@ class default_preconditioner(_feature_flag):
     Add a diagonal correction to scalable inducing point methods
     """
 
-    pass
-
 
 __all__ = ["checkpoint_kernel", "default_preconditioner"]

gpytorch/constraints/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .constraints import GreaterThan, Interval, LessThan, Positive
 
 __all__ = ["GreaterThan", "Interval", "LessThan", "Positive"]

gpytorch/constraints/constraints.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import math
-from typing import Optional
 
 import torch
 from torch import sigmoid, Tensor
@@ -87,7 +86,7 @@ def check(self, tensor) -> bool:
 
     def check_raw(self, tensor) -> bool:
         return bool(
-            torch.all((self.transform(tensor) <= self.upper_bound))
+            torch.all(self.transform(tensor) <= self.upper_bound)
             and torch.all(self.transform(tensor) >= self.lower_bound)
         )
 
@@ -137,7 +136,7 @@ def inverse_transform(self, transformed_tensor: Tensor) -> Tensor:
         return tensor
 
     @property
-    def initial_value(self) -> Optional[Tensor]:
+    def initial_value(self) -> Tensor | None:
         """
         The initial parameter value (if specified, None otherwise)
         """

gpytorch/distributions/__init__.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 from .delta import Delta
 from .distribution import Distribution
 from .multitask_multivariate_normal import MultitaskMultivariateNormal

gpytorch/distributions/delta.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import numbers
 
 import torch
@@ -34,14 +36,14 @@ class Delta(Distribution):
 
         def __init__(self, v, log_density=0.0, event_dim=0, validate_args=None):
             if event_dim > v.dim():
-                raise ValueError("Expected event_dim <= v.dim(), actual {} vs {}".format(event_dim, v.dim()))
+                raise ValueError(f"Expected event_dim <= v.dim(), actual {event_dim} vs {v.dim()}")
             batch_dim = v.dim() - event_dim
             batch_shape = v.shape[:batch_dim]
             event_shape = v.shape[batch_dim:]
             if isinstance(log_density, numbers.Number):
                 log_density = torch.full(batch_shape, log_density, dtype=v.dtype, device=v.device)
             elif validate_args and log_density.shape != batch_shape:
-                raise ValueError("Expected log_density.shape = {}, actual {}".format(log_density.shape, batch_shape))
+                raise ValueError(f"Expected log_density.shape = {log_density.shape}, actual {batch_shape}")
             self.v = v
             self.log_density = log_density
             super().__init__(batch_shape, event_shape, validate_args=validate_args)

gpytorch/distributions/distribution.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 from torch.distributions import Distribution as TDistribution
 
 

gpytorch/distributions/multitask_multivariate_normal.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import torch
 from linear_operator import LinearOperator, to_linear_operator
 from linear_operator.operators import (

gpytorch/distributions/multivariate_normal.py

@@ -5,7 +5,6 @@
 import math
 import warnings
 from numbers import Number
-from typing import Optional, Tuple, Union
 
 import torch
 from linear_operator import to_dense, to_linear_operator
@@ -42,7 +41,7 @@ class MultivariateNormal(TMultivariateNormal, Distribution):
     :ivar torch.Tensor variance: The variance.
     """
 
-    def __init__(self, mean: Tensor, covariance_matrix: Union[Tensor, LinearOperator], validate_args: bool = False):
+    def __init__(self, mean: Tensor, covariance_matrix: Tensor | LinearOperator, validate_args: bool = False):
         self._islazy = isinstance(mean, LinearOperator) or isinstance(covariance_matrix, LinearOperator)
         if self._islazy:
             if validate_args:
@@ -78,7 +77,7 @@ def _extended_shape(self, sample_shape: torch.Size = torch.Size()) -> torch.Size
         return sample_shape + self._batch_shape + self.base_sample_shape
 
     @staticmethod
-    def _repr_sizes(mean: Tensor, covariance_matrix: Union[Tensor, LinearOperator]) -> str:
+    def _repr_sizes(mean: Tensor, covariance_matrix: Tensor | LinearOperator) -> str:
         return f"MultivariateNormal(loc: {mean.size()}, scale: {covariance_matrix.size()})"
 
     @property
@@ -119,7 +118,7 @@ def covariance_matrix(self) -> Tensor:
         else:
             return super().covariance_matrix
 
-    def confidence_region(self) -> Tuple[Tensor, Tensor]:
+    def confidence_region(self) -> tuple[Tensor, Tensor]:
         """
         Returns 2 standard deviations above and below the mean.
 
@@ -252,7 +251,7 @@ def log_prob(self, value: Tensor) -> Tensor:
         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])
         return res
 
-    def rsample(self, sample_shape: torch.Size = torch.Size(), base_samples: Optional[Tensor] = None) -> Tensor:
+    def rsample(self, sample_shape: torch.Size = torch.Size(), base_samples: Tensor | None = None) -> Tensor:
         r"""
         Generates a `sample_shape` shaped reparameterized sample or `sample_shape`
         shaped batch of reparameterized samples if the distribution parameters
@@ -320,7 +319,7 @@ def rsample(self, sample_shape: torch.Size = torch.Size(), base_samples: Optiona
 
         return res
 
-    def sample(self, sample_shape: torch.Size = torch.Size(), base_samples: Optional[Tensor] = None) -> Tensor:
+    def sample(self, sample_shape: torch.Size = torch.Size(), base_samples: Tensor | None = None) -> Tensor:
         r"""
         Generates a `sample_shape` shaped sample or `sample_shape`
         shaped batch of samples if the distribution parameters
@@ -391,7 +390,7 @@ def __add__(self, other: MultivariateNormal) -> MultivariateNormal:
         elif isinstance(other, int) or isinstance(other, float):
             return self.__class__(self.mean + other, self.lazy_covariance_matrix)
         else:
-            raise RuntimeError("Unsupported type {} for addition w/ MultivariateNormal".format(type(other)))
+            raise RuntimeError(f"Unsupported type {type(other)} for addition w/ MultivariateNormal")
 
     def __getitem__(self, idx) -> MultivariateNormal:
         r"""