Replace list[SupervisedDataset] hints with Sequence[SupervisedDataset] in Generator.fit & cross_validate

ax/models/torch/botorch.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from copy import deepcopy
 from logging import Logger
 from typing import Any, Optional
@@ -292,7 +292,7 @@ def __init__(
     @copy_doc(TorchGenerator.fit)
     def fit(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         search_space_digest: SearchSpaceDigest,
         candidate_metadata: list[list[TCandidateMetadata]] | None = None,
     ) -> None:
@@ -470,7 +470,7 @@ def best_point(
     @copy_doc(TorchGenerator.cross_validate)
     def cross_validate(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         X_test: Tensor,
         search_space_digest: SearchSpaceDigest,
         use_posterior_predictive: bool = False,

ax/models/torch/cbo_lcea.py

@@ -6,6 +6,7 @@
 
 # pyre-strict
 
+from collections.abc import Sequence
 from logging import Logger
 from typing import Any, cast, Union
 
@@ -117,7 +118,7 @@ def __init__(
     @copy_doc(TorchGenerator.fit)
     def fit(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         search_space_digest: SearchSpaceDigest,
         candidate_metadata: list[list[TCandidateMetadata]] | None = None,
     ) -> None:

ax/models/torch/cbo_sac.py

@@ -6,6 +6,7 @@
 
 # pyre-strict
 
+from collections.abc import Sequence
 from logging import Logger
 from typing import Any
 
@@ -51,7 +52,7 @@ def __init__(self, decomposition: dict[str, list[str]]) -> None:
     @copy_doc(TorchGenerator.fit)
     def fit(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         search_space_digest: SearchSpaceDigest,
         candidate_metadata: list[list[TCandidateMetadata]] | None = None,
     ) -> None:

ax/models/torch/randomforest.py

@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 import numpy as np
 import numpy.typing as npt
 import torch
@@ -45,7 +47,7 @@ def __init__(self, max_features: str | None = "sqrt", num_trees: int = 500) -> N
     @copy_doc(TorchGenerator.fit)
     def fit(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         search_space_digest: SearchSpaceDigest,
         candidate_metadata: list[list[TCandidateMetadata]] | None = None,
     ) -> None:
@@ -68,7 +70,7 @@ def predict(self, X: Tensor) -> tuple[Tensor, Tensor]:
     @copy_doc(TorchGenerator.cross_validate)
     def cross_validate(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         X_test: Tensor,
         search_space_digest: SearchSpaceDigest,
         use_posterior_predictive: bool = False,

ax/models/torch/utils.py

@@ -6,7 +6,7 @@
 
 # pyre-strict
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from logging import Logger
 from typing import Any, cast
@@ -586,7 +586,7 @@ def randomize_objective_weights(
 
 
 def _datasets_to_legacy_inputs(
-    datasets: list[SupervisedDataset],
+    datasets: Sequence[SupervisedDataset],
 ) -> tuple[list[Tensor], list[Tensor], list[Tensor]]:
     """Convert a dictionary of dataset containers to legacy X, Y, Yvar inputs"""
     Xs, Ys, Yvars = [], [], []

ax/models/torch_base.py

@@ -9,7 +9,6 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping, Sequence
-
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -123,7 +122,7 @@ class TorchGenerator(BaseGenerator):
 
     def fit(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         search_space_digest: SearchSpaceDigest,
         candidate_metadata: list[list[TCandidateMetadata]] | None = None,
     ) -> None:
@@ -199,7 +198,7 @@ def best_point(
 
     def cross_validate(
         self,
-        datasets: list[SupervisedDataset],
+        datasets: Sequence[SupervisedDataset],
         X_test: Tensor,
         search_space_digest: SearchSpaceDigest,
         use_posterior_predictive: bool = False,