Apply PEP 604 union type syntax codemod

I came across https://github.com/asottile/pyupgrade randomly and figure I give it a shot.

This codemods all `Optional[X]` type definitions to use the PEP 604 syntax `X | None` instead. I also ran ufmt to fix the imports.

botorch/acquisition/acquisition.py

@@ -41,7 +41,7 @@ def __init__(self, model: Model) -> None:
         super().__init__()
         self.model: Model = model
 
-    def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
+    def set_X_pending(self, X_pending: Tensor | None = None) -> None:
         r"""Informs the acquisition function about pending design points.
 
         Args:
@@ -115,7 +115,7 @@ class MCSamplerMixin(ABC):
 
     _default_sample_shape = torch.Size([512])
 
-    def __init__(self, sampler: Optional[MCSampler] = None) -> None:
+    def __init__(self, sampler: MCSampler | None = None) -> None:
         r"""Register the sampler on the acquisition function.
 
         Args:

botorch/acquisition/active_learning.py

@@ -53,9 +53,9 @@ def __init__(
         self,
         model: Model,
         mc_points: Tensor,
-        sampler: Optional[MCSampler] = None,
-        posterior_transform: Optional[PosteriorTransform] = None,
-        X_pending: Optional[Tensor] = None,
+        sampler: MCSampler | None = None,
+        posterior_transform: PosteriorTransform | None = None,
+        X_pending: Tensor | None = None,
     ) -> None:
         r"""q-Integrated Negative Posterior Variance.
 
@@ -140,7 +140,7 @@ def __init__(
         self,
         model: Model,
         objective: MCAcquisitionObjective,
-        sampler: Optional[MCSampler] = None,
+        sampler: MCSampler | None = None,
     ) -> None:
         r"""Pairwise Monte Carlo Posterior Variance
 

botorch/acquisition/analytic.py

@@ -52,7 +52,7 @@ class AnalyticAcquisitionFunction(AcquisitionFunction, ABC):
     def __init__(
         self,
         model: Model,
-        posterior_transform: Optional[PosteriorTransform] = None,
+        posterior_transform: PosteriorTransform | None = None,
     ) -> None:
         r"""Base constructor for analytic acquisition functions.
 
@@ -76,14 +76,14 @@ def __init__(
                 )
         self.posterior_transform = posterior_transform
 
-    def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
+    def set_X_pending(self, X_pending: Tensor | None = None) -> None:
         raise UnsupportedError(
             "Analytic acquisition functions do not account for X_pending yet."
         )
 
     def _mean_and_sigma(
         self, X: Tensor, compute_sigma: bool = True, min_var: float = 1e-12
-    ) -> tuple[Tensor, Optional[Tensor]]:
+    ) -> tuple[Tensor, Tensor | None]:
         """Computes the first and second moments of the model posterior.
 
         Args:
@@ -135,8 +135,8 @@ class LogProbabilityOfImprovement(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        best_f: Union[float, Tensor],
-        posterior_transform: Optional[PosteriorTransform] = None,
+        best_f: float | Tensor,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ):
         r"""Single-outcome Probability of Improvement.
@@ -189,8 +189,8 @@ class ProbabilityOfImprovement(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        best_f: Union[float, Tensor],
-        posterior_transform: Optional[PosteriorTransform] = None,
+        best_f: float | Tensor,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ):
         r"""Single-outcome Probability of Improvement.
@@ -237,8 +237,8 @@ class qAnalyticProbabilityOfImprovement(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        best_f: Union[float, Tensor],
-        posterior_transform: Optional[PosteriorTransform] = None,
+        best_f: float | Tensor,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ) -> None:
         """qPI using an analytic approximation.
@@ -314,8 +314,8 @@ class ExpectedImprovement(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        best_f: Union[float, Tensor],
-        posterior_transform: Optional[PosteriorTransform] = None,
+        best_f: float | Tensor,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ):
         r"""Single-outcome Expected Improvement (analytic).
@@ -378,8 +378,8 @@ class LogExpectedImprovement(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        best_f: Union[float, Tensor],
-        posterior_transform: Optional[PosteriorTransform] = None,
+        best_f: float | Tensor,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ):
         r"""Logarithm of single-outcome Expected Improvement (analytic).
@@ -447,9 +447,9 @@ class LogConstrainedExpectedImprovement(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        best_f: Union[float, Tensor],
+        best_f: float | Tensor,
         objective_index: int,
-        constraints: dict[int, tuple[Optional[float], Optional[float]]],
+        constraints: dict[int, tuple[float | None, float | None]],
         maximize: bool = True,
     ) -> None:
         r"""Analytic Log Constrained Expected Improvement.
@@ -525,9 +525,9 @@ class ConstrainedExpectedImprovement(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        best_f: Union[float, Tensor],
+        best_f: float | Tensor,
         objective_index: int,
-        constraints: dict[int, tuple[Optional[float], Optional[float]]],
+        constraints: dict[int, tuple[float | None, float | None]],
         maximize: bool = True,
     ) -> None:
         r"""Analytic Constrained Expected Improvement.
@@ -606,7 +606,7 @@ def __init__(
         X_observed: Tensor,
         num_fantasies: int = 20,
         maximize: bool = True,
-        posterior_transform: Optional[PosteriorTransform] = None,
+        posterior_transform: PosteriorTransform | None = None,
     ) -> None:
         r"""Single-outcome Noisy Log Expected Improvement (via fantasies).
 
@@ -762,8 +762,8 @@ class UpperConfidenceBound(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        beta: Union[float, Tensor],
-        posterior_transform: Optional[PosteriorTransform] = None,
+        beta: float | Tensor,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ) -> None:
         r"""Single-outcome Upper Confidence Bound.
@@ -812,7 +812,7 @@ class PosteriorMean(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        posterior_transform: Optional[PosteriorTransform] = None,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ) -> None:
         r"""Single-outcome Posterior Mean.
@@ -857,7 +857,7 @@ def __init__(
         self,
         model: Model,
         weights: Tensor,
-        posterior_transform: Optional[PosteriorTransform] = None,
+        posterior_transform: PosteriorTransform | None = None,
     ) -> None:
         r"""Scalarized Posterior Mean.
 
@@ -919,7 +919,7 @@ class PosteriorStandardDeviation(AnalyticAcquisitionFunction):
     def __init__(
         self,
         model: Model,
-        posterior_transform: Optional[PosteriorTransform] = None,
+        posterior_transform: PosteriorTransform | None = None,
         maximize: bool = True,
     ) -> None:
         r"""Single-outcome Posterior Mean.
@@ -1135,8 +1135,8 @@ def _get_noiseless_fantasy_model(
 
 
 def _preprocess_constraint_bounds(
-    acqf: Union[LogConstrainedExpectedImprovement, ConstrainedExpectedImprovement],
-    constraints: dict[int, tuple[Optional[float], Optional[float]]],
+    acqf: LogConstrainedExpectedImprovement | ConstrainedExpectedImprovement,
+    constraints: dict[int, tuple[float | None, float | None]],
 ) -> None:
     r"""Set up constraint bounds.
 
@@ -1180,7 +1180,7 @@ def _preprocess_constraint_bounds(
 
 
 def _compute_log_prob_feas(
-    acqf: Union[LogConstrainedExpectedImprovement, ConstrainedExpectedImprovement],
+    acqf: LogConstrainedExpectedImprovement | ConstrainedExpectedImprovement,
     means: Tensor,
     sigmas: Tensor,
 ) -> Tensor:

botorch/acquisition/bayesian_active_learning.py

@@ -79,10 +79,10 @@ class qBayesianActiveLearningByDisagreement(
 ):
     def __init__(
         self,
-        model: Union[ModelListGP, SaasFullyBayesianSingleTaskGP],
-        sampler: Optional[MCSampler] = None,
-        posterior_transform: Optional[PosteriorTransform] = None,
-        X_pending: Optional[Tensor] = None,
+        model: ModelListGP | SaasFullyBayesianSingleTaskGP,
+        sampler: MCSampler | None = None,
+        posterior_transform: PosteriorTransform | None = None,
+        X_pending: Tensor | None = None,
     ) -> None:
         """
         Batch implementation [kirsch2019batchbald]_ of BALD [Houlsby2011bald]_,

botorch/acquisition/cached_cholesky.py

@@ -72,7 +72,7 @@ def __init__(
         self,
         model: Model,
         cache_root: bool = False,
-        sampler: Optional[MCSampler] = None,
+        sampler: MCSampler | None = None,
     ) -> None:
         r"""Set class attributes and perform compatibility checks.
 

botorch/acquisition/cost_aware.py

@@ -13,7 +13,8 @@
 
 import warnings
 from abc import ABC, abstractmethod
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 
 import torch
 from botorch import settings
@@ -35,7 +36,7 @@ class CostAwareUtility(Module, ABC):
 
     @abstractmethod
     def forward(
-        self, X: Tensor, deltas: Tensor, sampler: Optional[MCSampler] = None
+        self, X: Tensor, deltas: Tensor, sampler: MCSampler | None = None
     ) -> Tensor:
         r"""Evaluate the cost-aware utility on the candidates and improvements.
 
@@ -67,7 +68,7 @@ def __init__(self, cost: Callable[[Tensor, Tensor], Tensor]) -> None:
         self._cost_callable: Callable[[Tensor, Tensor], Tensor] = cost
 
     def forward(
-        self, X: Tensor, deltas: Tensor, sampler: Optional[MCSampler] = None
+        self, X: Tensor, deltas: Tensor, sampler: MCSampler | None = None
     ) -> Tensor:
         r"""Evaluate the cost function on the candidates and improvements.
 
@@ -109,9 +110,9 @@ class InverseCostWeightedUtility(CostAwareUtility):
 
     def __init__(
         self,
-        cost_model: Union[DeterministicModel, GPyTorchModel],
+        cost_model: DeterministicModel | GPyTorchModel,
         use_mean: bool = True,
-        cost_objective: Optional[MCAcquisitionObjective] = None,
+        cost_objective: MCAcquisitionObjective | None = None,
         min_cost: float = 1e-2,
     ) -> None:
         r"""Cost-aware utility that weights increase in utility by inverse cost.
@@ -153,8 +154,8 @@ def forward(
         self,
         X: Tensor,
         deltas: Tensor,
-        sampler: Optional[MCSampler] = None,
-        X_evaluation_mask: Optional[Tensor] = None,
+        sampler: MCSampler | None = None,
+        X_evaluation_mask: Tensor | None = None,
     ) -> Tensor:
         r"""Evaluate the cost function on the candidates and improvements. Note
         that negative values of `deltas` are instead scaled by the cost, and not

botorch/acquisition/decoupled.py

@@ -52,7 +52,7 @@ class DecoupledAcquisitionFunction(AcquisitionFunction, ABC):
     """
 
     def __init__(
-        self, model: ModelList, X_evaluation_mask: Optional[Tensor] = None, **kwargs
+        self, model: ModelList, X_evaluation_mask: Tensor | None = None, **kwargs
     ) -> None:
         r"""Initialize.
 
@@ -71,12 +71,12 @@ def __init__(
         self.X_pending = None
 
     @property
-    def X_evaluation_mask(self) -> Optional[Tensor]:
+    def X_evaluation_mask(self) -> Tensor | None:
         r"""Get the evaluation indices for the new candidate."""
         return self._X_evaluation_mask
 
     @X_evaluation_mask.setter
-    def X_evaluation_mask(self, X_evaluation_mask: Optional[Tensor] = None) -> None:
+    def X_evaluation_mask(self, X_evaluation_mask: Tensor | None = None) -> None:
         r"""Set the evaluation indices for the new candidate."""
         if X_evaluation_mask is not None:
             # TODO: Add batch support
@@ -92,8 +92,8 @@ def X_evaluation_mask(self, X_evaluation_mask: Optional[Tensor] = None) -> None:
 
     def set_X_pending(
         self,
-        X_pending: Optional[Tensor] = None,
-        X_pending_evaluation_mask: Optional[Tensor] = None,
+        X_pending: Tensor | None = None,
+        X_pending_evaluation_mask: Tensor | None = None,
     ) -> None:
         r"""Informs the AF about pending design points for different outcomes.
 
@@ -135,7 +135,7 @@ def set_X_pending(
             self.X_pending = X_pending
             self.X_pending_evaluation_mask = X_pending_evaluation_mask
 
-    def construct_evaluation_mask(self, X: Tensor) -> Optional[Tensor]:
+    def construct_evaluation_mask(self, X: Tensor) -> Tensor | None:
         r"""Construct the boolean evaluation mask for X and X_pending
 
         Args:

botorch/acquisition/factory.py

@@ -10,7 +10,9 @@
 
 from __future__ import annotations
 
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+
+from typing import Optional, Union
 
 import torch
 
@@ -35,21 +37,21 @@ def get_acquisition_function(
     model: Model,
     objective: MCAcquisitionObjective,
     X_observed: Tensor,
-    posterior_transform: Optional[PosteriorTransform] = None,
-    X_pending: Optional[Tensor] = None,
-    constraints: Optional[list[Callable[[Tensor], Tensor]]] = None,
-    eta: Optional[Union[Tensor, float]] = 1e-3,
+    posterior_transform: PosteriorTransform | None = None,
+    X_pending: Tensor | None = None,
+    constraints: list[Callable[[Tensor], Tensor]] | None = None,
+    eta: Tensor | float | None = 1e-3,
     mc_samples: int = 512,
-    seed: Optional[int] = None,
+    seed: int | None = None,
     *,
     # optional parameters that are only needed for certain acquisition functions
     tau: float = 1e-3,
     prune_baseline: bool = True,
-    marginalize_dim: Optional[int] = None,
+    marginalize_dim: int | None = None,
     cache_root: bool = True,
-    beta: Optional[float] = None,
-    ref_point: Union[None, list[float], Tensor] = None,
-    Y: Optional[Tensor] = None,
+    beta: float | None = None,
+    ref_point: None | list[float] | Tensor = None,
+    Y: Tensor | None = None,
     alpha: float = 0.0,
 ) -> monte_carlo.MCAcquisitionFunction:
     r"""Convenience function for initializing botorch acquisition functions.