From fee56323970cf10c127fe78fdefb0f0fc9e34dab Mon Sep 17 00:00:00 2001 From: Miles Olson Date: Tue, 29 Oct 2024 09:11:39 -0700 Subject: [PATCH] Fix np.ndarray type annotation (#2983) Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/2983 Reviewed By: paschai Differential Revision: D65146743 --- ax/benchmark/benchmark.py | 7 +- ax/benchmark/benchmark_result.py | 19 +-- ax/benchmark/problems/hpo/torchvision.py | 1 - ax/benchmark/runners/base.py | 6 +- .../runners/test_botorch_test_problem.py | 2 - ax/core/batch_trial.py | 1 - ax/core/experiment.py | 1 - ax/core/observation.py | 8 +- ax/core/parameter.py | 1 - ax/core/search_space.py | 8 +- ax/core/utils.py | 8 +- ax/early_stopping/strategies/base.py | 11 +- ax/metrics/branin.py | 11 +- ax/metrics/branin_map.py | 7 +- ax/metrics/hartmann6.py | 8 +- ax/metrics/l2norm.py | 4 +- ax/metrics/noisy_function.py | 4 +- ax/metrics/noisy_function_map.py | 4 +- ax/metrics/sklearn.py | 5 +- ax/modelbridge/best_model_selector.py | 13 +- ax/modelbridge/cross_validation.py | 25 ++-- ax/modelbridge/discrete.py | 4 - ax/modelbridge/map_torch.py | 5 +- ax/modelbridge/modelbridge_utils.py | 68 ++++------ ax/modelbridge/random.py | 2 - ax/modelbridge/tests/test_cross_validation.py | 8 -- .../tests/test_modelbridge_utils.py | 3 +- .../tests/test_torch_modelbridge.py | 4 +- ax/modelbridge/torch.py | 16 +-- ax/modelbridge/transforms/choice_encode.py | 4 +- ax/modelbridge/transforms/ivw.py | 8 +- ax/modelbridge/transforms/log_y.py | 36 ++---- ax/modelbridge/transforms/metrics_as_task.py | 2 +- ax/modelbridge/transforms/relativize.py | 16 +-- ax/modelbridge/transforms/rounding.py | 7 +- .../transforms/tests/test_log_y_transform.py | 2 - .../tests/test_relativize_transform.py | 7 +- .../tests/test_transform_to_new_sq.py | 4 +- .../transforms/transform_to_new_sq.py | 13 +- ax/modelbridge/transforms/winsorize.py | 4 +- ax/models/discrete/full_factorial.py | 8 +- ax/models/discrete/thompson.py | 28 ++-- ax/models/discrete_base.py | 20 +-- ax/models/model_utils.py | 60 +++------ ax/models/random/base.py | 26 ++-- ax/models/random/sobol.py | 16 +-- ax/models/random/uniform.py | 7 +- ax/models/torch/botorch.py | 8 +- ax/models/torch/botorch_modular/model.py | 5 +- ax/models/torch/randomforest.py | 10 +- ax/models/torch/tests/test_acquisition.py | 2 - ax/models/torch/tests/test_sebo.py | 2 - ax/models/torch/utils.py | 7 +- ax/plot/contour.py | 10 +- ax/plot/feature_importances.py | 7 +- ax/plot/helper.py | 4 +- ax/plot/pareto_frontier.py | 16 +-- ax/plot/pareto_utils.py | 10 +- ax/plot/slice.py | 4 +- ax/plot/trace.py | 34 ++--- ax/service/tests/scheduler_test_utils.py | 1 - ax/service/tests/test_instantiation_utils.py | 14 -- ax/service/tests/test_managed_loop.py | 11 +- ax/service/utils/report_utils.py | 10 +- ax/utils/measurement/synthetic_functions.py | 40 ++---- .../tests/test_synthetic_functions.py | 2 - ax/utils/sensitivity/sobol_measures.py | 12 +- .../sensitivity/tests/test_sensitivity.py | 13 -- ax/utils/stats/model_fit_stats.py | 122 ++++++------------ ax/utils/stats/statstools.py | 52 +++----- ax/utils/testing/mock.py | 4 +- 71 files changed, 324 insertions(+), 608 deletions(-) diff --git a/ax/benchmark/benchmark.py b/ax/benchmark/benchmark.py index 1375ff76c43..352f1768b76 100644 --- a/ax/benchmark/benchmark.py +++ b/ax/benchmark/benchmark.py @@ -25,6 +25,7 @@ from time import monotonic, time import numpy as np +import numpy.typing as npt from ax.benchmark.benchmark_method import BenchmarkMethod from ax.benchmark.benchmark_problem import BenchmarkProblem @@ -41,12 +42,10 @@ def compute_score_trace( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - optimization_trace: np.ndarray, + optimization_trace: npt.NDArray, num_baseline_trials: int, problem: BenchmarkProblem, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: +) -> npt.NDArray: """Computes a score trace from the optimization trace.""" # Use the first GenerationStep's best found point as baseline. Sometimes (ex. in diff --git a/ax/benchmark/benchmark_result.py b/ax/benchmark/benchmark_result.py index 10acf37ec95..2811769cf52 100644 --- a/ax/benchmark/benchmark_result.py +++ b/ax/benchmark/benchmark_result.py @@ -13,10 +13,10 @@ from collections.abc import Iterable from dataclasses import dataclass -import numpy as np +import numpy.typing as npt from ax.core.experiment import Experiment from ax.utils.common.base import Base -from numpy import nanmean, nanquantile, ndarray +from numpy import nanmean, nanquantile from pandas import DataFrame from scipy.stats import sem @@ -78,14 +78,10 @@ class BenchmarkResult(Base): name: str seed: int - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. - oracle_trace: ndarray - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. - inference_trace: ndarray - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. - optimization_trace: ndarray - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. - score_trace: ndarray + oracle_trace: npt.NDArray + inference_trace: npt.NDArray + optimization_trace: npt.NDArray + score_trace: npt.NDArray fit_time: float gen_time: float @@ -160,8 +156,7 @@ def from_benchmark_results( def _get_stats( - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. - step_data: Iterable[np.ndarray], + step_data: Iterable[npt.NDArray], percentiles: list[float], ) -> dict[str, list[float]]: quantiles = [] diff --git a/ax/benchmark/problems/hpo/torchvision.py b/ax/benchmark/problems/hpo/torchvision.py index 654eb723d64..1a6e424a5f2 100644 --- a/ax/benchmark/problems/hpo/torchvision.py +++ b/ax/benchmark/problems/hpo/torchvision.py @@ -111,7 +111,6 @@ def train_and_evaluate( total += labels.size(0) correct += (predicted == labels).sum().item() - # pyre-fixme[7]: Expected `Tensor` but got `float`. return correct / total diff --git a/ax/benchmark/runners/base.py b/ax/benchmark/runners/base.py index b00f324bd27..2415893be7d 100644 --- a/ax/benchmark/runners/base.py +++ b/ax/benchmark/runners/base.py @@ -11,6 +11,8 @@ from math import sqrt from typing import Any +import numpy.typing as npt + import torch from ax.core.base_trial import BaseTrial, TrialStatus from ax.core.batch_trial import BatchTrial @@ -22,7 +24,6 @@ from ax.utils.common.serialization import TClassDecoderRegistry, TDecoderRegistry from ax.utils.common.typeutils import checked_cast -from numpy import ndarray from torch import Tensor @@ -68,8 +69,7 @@ def get_Y_true(self, params: Mapping[str, TParamValue]) -> Tensor: """ ... - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. - def evaluate_oracle(self, parameters: Mapping[str, TParamValue]) -> ndarray: + def evaluate_oracle(self, parameters: Mapping[str, TParamValue]) -> npt.NDArray: """ Evaluate oracle metric values at a parameterization. In the base class, oracle values are underlying noiseless function values evaluated at the diff --git a/ax/benchmark/tests/runners/test_botorch_test_problem.py b/ax/benchmark/tests/runners/test_botorch_test_problem.py index c32c9a30b7c..9dcba5c4872 100644 --- a/ax/benchmark/tests/runners/test_botorch_test_problem.py +++ b/ax/benchmark/tests/runners/test_botorch_test_problem.py @@ -263,8 +263,6 @@ def test_synthetic_runner(self) -> None: nullcontext() if not isinstance(test_problem, SurrogateTestFunction) else patch.object( - # pyre-fixme: ParamBasedTestProblem` has no attribute - # `_surrogate`. runner.test_problem._surrogate, "predict", return_value=({"branin": [4.2]}, None), diff --git a/ax/core/batch_trial.py b/ax/core/batch_trial.py index 266897cb14f..fbc36778ce5 100644 --- a/ax/core/batch_trial.py +++ b/ax/core/batch_trial.py @@ -514,7 +514,6 @@ def normalized_arm_weights( weights = np.array(self.weights) if trunc_digits is not None: atomic_weight = 10**-trunc_digits - # pyre-fixme[16]: `float` has no attribute `astype`. int_weights = ( (total / atomic_weight) * (weights / np.sum(weights)) ).astype(int) diff --git a/ax/core/experiment.py b/ax/core/experiment.py index 4aaeab1a649..6d33a26915d 100644 --- a/ax/core/experiment.py +++ b/ax/core/experiment.py @@ -69,7 +69,6 @@ } -# pyre-fixme[13]: Attribute `_search_space` is never initialized. class Experiment(Base): """Base class for defining an experiment.""" diff --git a/ax/core/observation.py b/ax/core/observation.py index d9e88458593..7dd6a7bd768 100644 --- a/ax/core/observation.py +++ b/ax/core/observation.py @@ -16,6 +16,7 @@ import ax.core.experiment as experiment import numpy as np +import numpy.typing as npt import pandas as pd from ax.core.arm import Arm from ax.core.base_trial import NON_ABANDONED_STATUSES, TrialStatus @@ -187,13 +188,10 @@ class ObservationData(Base): """ def __init__( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. self, metric_names: list[str], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - means: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - covariance: np.ndarray, + means: npt.NDArray, + covariance: npt.NDArray, ) -> None: k = len(metric_names) if means.shape != (k,): diff --git a/ax/core/parameter.py b/ax/core/parameter.py index 10ead1e4846..9ee8445ef75 100644 --- a/ax/core/parameter.py +++ b/ax/core/parameter.py @@ -153,7 +153,6 @@ def dependents(self) -> dict[TParamValue, list[str]]: # pyre-fixme[7]: Expected `Parameter` but got implicit return value of `None`. def clone(self) -> Parameter: - # pyre-fixme[7]: Expected `Parameter` but got implicit return value of `None`. pass @property diff --git a/ax/core/search_space.py b/ax/core/search_space.py index d31b1bda78c..b77cfd8dff2 100644 --- a/ax/core/search_space.py +++ b/ax/core/search_space.py @@ -17,7 +17,7 @@ from random import choice, uniform from typing import Sequence -import numpy as np +import numpy.typing as npt import pandas as pd from ax import core from ax.core.arm import Arm @@ -1103,10 +1103,8 @@ class RobustSearchSpaceDigest: Only relevant if paired with a `distribution_sampler`. """ - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sample_param_perturbations: Callable[[], np.ndarray] | None = None - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sample_environmental: Callable[[], np.ndarray] | None = None + sample_param_perturbations: Callable[[], npt.NDArray] | None = None + sample_environmental: Callable[[], npt.NDArray] | None = None environmental_variables: list[str] = field(default_factory=list) multiplicative: bool = False diff --git a/ax/core/utils.py b/ax/core/utils.py index 33b41fe8773..33fa0d027fe 100644 --- a/ax/core/utils.py +++ b/ax/core/utils.py @@ -12,6 +12,7 @@ from typing import NamedTuple import numpy as np +import numpy.typing as npt from ax.core.arm import Arm from ax.core.base_trial import BaseTrial, TrialStatus from ax.core.batch_trial import BatchTrial @@ -129,12 +130,9 @@ def _get_missing_arm_trial_pairs(data: Data, metric_name: str) -> set[TArmTrial] def best_feasible_objective( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. optimization_config: OptimizationConfig, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - values: dict[str, np.ndarray], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: + values: dict[str, npt.NDArray], +) -> npt.NDArray: """Compute the best feasible objective value found by each iteration. Args: diff --git a/ax/early_stopping/strategies/base.py b/ax/early_stopping/strategies/base.py index 6096c9e3ecf..a5f073f9158 100644 --- a/ax/early_stopping/strategies/base.py +++ b/ax/early_stopping/strategies/base.py @@ -12,7 +12,7 @@ from dataclasses import dataclass from logging import Logger -import numpy as np +import numpy.typing as npt import pandas as pd from ax.core.base_trial import TrialStatus from ax.core.data import Data @@ -54,12 +54,9 @@ class EarlyStoppingTrainingData: which data come from the same arm. """ - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - X: np.ndarray - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Y: np.ndarray - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Yvar: np.ndarray + X: npt.NDArray + Y: npt.NDArray + Yvar: npt.NDArray arm_names: list[str | None] diff --git a/ax/metrics/branin.py b/ax/metrics/branin.py index 73e39f33763..e1d97b9336f 100644 --- a/ax/metrics/branin.py +++ b/ax/metrics/branin.py @@ -6,27 +6,24 @@ # pyre-strict -import numpy as np +import numpy.typing as npt from ax.metrics.noisy_function import NoisyFunctionMetric from ax.utils.common.typeutils import checked_cast from ax.utils.measurement.synthetic_functions import aug_branin, branin class BraninMetric(NoisyFunctionMetric): - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> float: + def f(self, x: npt.NDArray) -> float: x1, x2 = x return checked_cast(float, branin(x1=x1, x2=x2)) class NegativeBraninMetric(BraninMetric): - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> float: + def f(self, x: npt.NDArray) -> float: fpos = super().f(x) return -fpos class AugmentedBraninMetric(NoisyFunctionMetric): - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> float: + def f(self, x: npt.NDArray) -> float: return checked_cast(float, aug_branin(x)) diff --git a/ax/metrics/branin_map.py b/ax/metrics/branin_map.py index 58a35e20372..686e3bca353 100644 --- a/ax/metrics/branin_map.py +++ b/ax/metrics/branin_map.py @@ -15,6 +15,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import pandas as pd from ax.core.base_trial import BaseTrial from ax.core.map_data import MapData, MapKeyInfo @@ -117,8 +118,7 @@ def fetch_trial_data( # pyre-fixme[14]: `f` overrides method defined in `NoisyFunctionMapMetric` # inconsistently. - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray, timestamp: int) -> Mapping[str, Any]: + def f(self, x: npt.NDArray, timestamp: int) -> Mapping[str, Any]: x1, x2 = x if self.rate is not None: @@ -161,8 +161,7 @@ def fetch_trial_data( **kwargs, ) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> Mapping[str, Any]: + def f(self, x: npt.NDArray) -> Mapping[str, Any]: if self.index < len(FIDELITY): self.index += 1 diff --git a/ax/metrics/hartmann6.py b/ax/metrics/hartmann6.py index 53944f5ebbd..40485e02307 100644 --- a/ax/metrics/hartmann6.py +++ b/ax/metrics/hartmann6.py @@ -6,19 +6,17 @@ # pyre-strict -import numpy as np +import numpy.typing as npt from ax.metrics.noisy_function import NoisyFunctionMetric from ax.utils.common.typeutils import checked_cast from ax.utils.measurement.synthetic_functions import aug_hartmann6, hartmann6 class Hartmann6Metric(NoisyFunctionMetric): - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> float: + def f(self, x: npt.NDArray) -> float: return checked_cast(float, hartmann6(x)) class AugmentedHartmann6Metric(NoisyFunctionMetric): - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> float: + def f(self, x: npt.NDArray) -> float: return checked_cast(float, aug_hartmann6(x)) diff --git a/ax/metrics/l2norm.py b/ax/metrics/l2norm.py index b17c0ddc933..50dbc6397dd 100644 --- a/ax/metrics/l2norm.py +++ b/ax/metrics/l2norm.py @@ -7,10 +7,10 @@ # pyre-strict import numpy as np +import numpy.typing as npt from ax.metrics.noisy_function import NoisyFunctionMetric class L2NormMetric(NoisyFunctionMetric): - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> float: + def f(self, x: npt.NDArray) -> float: return np.sqrt((x**2).sum()) diff --git a/ax/metrics/noisy_function.py b/ax/metrics/noisy_function.py index f802c89fe57..d83e7497b2a 100644 --- a/ax/metrics/noisy_function.py +++ b/ax/metrics/noisy_function.py @@ -13,6 +13,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import pandas as pd from ax.core.base_trial import BaseTrial from ax.core.data import Data @@ -104,8 +105,7 @@ def _evaluate(self, params: TParameterization) -> float: x = np.array([params[p] for p in self.param_names]) return self.f(x) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> float: + def f(self, x: npt.NDArray) -> float: """The deterministic function that produces the metric outcomes.""" raise NotImplementedError diff --git a/ax/metrics/noisy_function_map.py b/ax/metrics/noisy_function_map.py index 32f2ac43efe..e0ea389a111 100644 --- a/ax/metrics/noisy_function_map.py +++ b/ax/metrics/noisy_function_map.py @@ -15,6 +15,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import pandas as pd from ax.core.base_trial import BaseTrial from ax.core.map_data import MapData, MapKeyInfo @@ -112,7 +113,6 @@ def fetch_trial_data( MetricFetchE(message=f"Failed to fetch {self.name}", exception=e) ) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, x: np.ndarray) -> Mapping[str, Any]: + def f(self, x: npt.NDArray) -> Mapping[str, Any]: """The deterministic function that produces the metric outcomes.""" raise NotImplementedError diff --git a/ax/metrics/sklearn.py b/ax/metrics/sklearn.py index 62d566583e2..06b04d04fef 100644 --- a/ax/metrics/sklearn.py +++ b/ax/metrics/sklearn.py @@ -14,7 +14,7 @@ from math import sqrt from typing import Any -import numpy as np +import numpy.typing as npt import pandas as pd from ax.core.arm import Arm from ax.core.base_trial import BaseTrial @@ -46,8 +46,7 @@ class SklearnDataset(Enum): @lru_cache(maxsize=8) # pyre-fixme[2]: Parameter must be annotated. -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _get_data(dataset) -> dict[str, np.ndarray]: +def _get_data(dataset) -> dict[str, npt.NDArray]: """Return sklearn dataset, loading and caching if necessary.""" if dataset is SklearnDataset.DIGITS: return datasets.load_digits() diff --git a/ax/modelbridge/best_model_selector.py b/ax/modelbridge/best_model_selector.py index 7829954c5c4..fef28a68e03 100644 --- a/ax/modelbridge/best_model_selector.py +++ b/ax/modelbridge/best_model_selector.py @@ -15,6 +15,7 @@ from typing import Any, Union import numpy as np +import numpy.typing as npt from ax.exceptions.core import UserInputError from ax.modelbridge.model_spec import ModelSpec from ax.utils.common.base import Base @@ -45,17 +46,13 @@ class ReductionCriterion(Enum): # NOTE: Callables need to be wrapped in `partial` to be registered as members. # pyre-fixme[35]: Target cannot be annotated. - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - MEAN: Callable[[ARRAYLIKE], np.ndarray] = partial(np.mean) + MEAN: Callable[[ARRAYLIKE], npt.NDArray] = partial(np.mean) # pyre-fixme[35]: Target cannot be annotated. - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - MIN: Callable[[ARRAYLIKE], np.ndarray] = partial(np.min) + MIN: Callable[[ARRAYLIKE], npt.NDArray] = partial(np.min) # pyre-fixme[35]: Target cannot be annotated. - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - MAX: Callable[[ARRAYLIKE], np.ndarray] = partial(np.max) + MAX: Callable[[ARRAYLIKE], npt.NDArray] = partial(np.max) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def __call__(self, array_like: ARRAYLIKE) -> np.ndarray: + def __call__(self, array_like: ARRAYLIKE) -> npt.NDArray: return self.value(array_like) diff --git a/ax/modelbridge/cross_validation.py b/ax/modelbridge/cross_validation.py index 63e8ae5fdbf..4755eca36a0 100644 --- a/ax/modelbridge/cross_validation.py +++ b/ax/modelbridge/cross_validation.py @@ -17,6 +17,7 @@ from warnings import warn import numpy as np +import numpy.typing as npt from ax.core.observation import Observation, ObservationData, recombine_observations from ax.core.optimization_config import OptimizationConfig from ax.modelbridge.base import ModelBridge, unwrap_observation_data @@ -328,8 +329,7 @@ def compute_diagnostics(result: list[CVResult]) -> CVDiagnostics: return diagnostics -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _arrayify_dict_values(d: dict[str, list[float]]) -> dict[str, np.ndarray]: +def _arrayify_dict_values(d: dict[str, list[float]]) -> dict[str, npt.NDArray]: """Helper to convert dictionary values to numpy arrays.""" return {k: np.array(v) for k, v in d.items()} @@ -403,10 +403,8 @@ def has_good_opt_config_model_fit( def _gen_train_test_split( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. folds: int, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - arm_names: np.ndarray, + arm_names: npt.NDArray, ) -> Iterable[tuple[set[str], set[str]]]: """Return train/test splits of arm names. @@ -539,8 +537,7 @@ def _model_fit_metric(metric_dict: dict[str, dict[str, float]]) -> float: return min(metric_dict["coefficient_of_determination"].values()) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _model_std_quality(std: np.ndarray) -> float: +def _model_std_quality(std: npt.NDArray) -> float: """Quantifies quality of the model uncertainty. A value of one means the uncertainty is perfectly predictive of the true standard deviation of the error. Values larger than one indicate over-estimation and negative values indicate @@ -566,11 +563,10 @@ def _model_std_quality(std: np.ndarray) -> float: def _predict_on_training_data( model_bridge: ModelBridge, untransform: bool = False, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. ) -> tuple[ - dict[str, np.ndarray], - dict[str, np.ndarray], - dict[str, np.ndarray], + dict[str, npt.NDArray], + dict[str, npt.NDArray], + dict[str, npt.NDArray], ]: """Makes predictions on the training data of a given experiment using a ModelBridge and returning the observed values, and the corresponding predictive means and @@ -631,11 +627,10 @@ def _predict_on_training_data( def _predict_on_cross_validation_data( model_bridge: ModelBridge, untransform: bool = False, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. ) -> tuple[ - dict[str, np.ndarray], - dict[str, np.ndarray], - dict[str, np.ndarray], + dict[str, npt.NDArray], + dict[str, npt.NDArray], + dict[str, npt.NDArray], ]: """Makes leave-one-out cross-validation predictions on the training data of the ModelBridge and returns the observed values, and the corresponding predictive means diff --git a/ax/modelbridge/discrete.py b/ax/modelbridge/discrete.py index bce00aa67db..e895d8378a1 100644 --- a/ax/modelbridge/discrete.py +++ b/ax/modelbridge/discrete.py @@ -35,10 +35,6 @@ FIT_MODEL_ERROR = "Model must be fit before {action}." -# pyre-fixme[13]: Attribute `model` is never initialized. -# pyre-fixme[13]: Attribute `outcomes` is never initialized. -# pyre-fixme[13]: Attribute `parameters` is never initialized. -# pyre-fixme[13]: Attribute `search_space` is never initialized. class DiscreteModelBridge(ModelBridge): """A model bridge for using models based on discrete parameters. diff --git a/ax/modelbridge/map_torch.py b/ax/modelbridge/map_torch.py index ee5b6a4877e..a3a6391767a 100644 --- a/ax/modelbridge/map_torch.py +++ b/ax/modelbridge/map_torch.py @@ -8,6 +8,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import torch from ax.core.base_trial import TrialStatus @@ -224,10 +225,8 @@ def _gen( ) def _array_to_observation_features( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. self, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - X: np.ndarray, + X: npt.NDArray, candidate_metadata: list[TCandidateMetadata] | None, ) -> list[ObservationFeatures]: """The difference b/t this method and diff --git a/ax/modelbridge/modelbridge_utils.py b/ax/modelbridge/modelbridge_utils.py index c9b47e3af76..cdea70c222d 100644 --- a/ax/modelbridge/modelbridge_utils.py +++ b/ax/modelbridge/modelbridge_utils.py @@ -16,6 +16,7 @@ from typing import Any, TYPE_CHECKING import numpy as np +import numpy.typing as npt import torch from ax.core.data import Data from ax.core.experiment import Experiment @@ -304,8 +305,7 @@ def extract_robust_digest( # NOTE: Extracting it from `param_names` in case the ordering is different. environmental_variables = param_names[num_non_env_vars:] - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def sample_environmental() -> np.ndarray: + def sample_environmental() -> npt.NDArray: """Get samples from the environmental distributions. Samples have the same dimension as the number of environmental variables. @@ -328,13 +328,11 @@ def sample_environmental() -> np.ndarray: environmental_variables = [] if len(pert_params) > 0: - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - constructor: Callable[[tuple[int, int]], np.ndarray] = ( + constructor: Callable[[tuple[int, int]], npt.NDArray] = ( np.ones if search_space.multiplicative else np.zeros ) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def sample_param_perturbations() -> np.ndarray: + def sample_param_perturbations() -> npt.NDArray: """Get samples of the input perturbations. Samples have the same dimension as the length of `param_names` @@ -368,8 +366,7 @@ def extract_objective_thresholds( objective_thresholds: TRefPoint, objective: Objective, outcomes: list[str], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray | None: +) -> npt.NDArray | None: """Extracts objective thresholds' values, in the order of `outcomes`. Will return None if no objective thresholds, otherwise the extracted tensor @@ -415,8 +412,7 @@ def extract_objective_thresholds( return obj_t -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def extract_objective_weights(objective: Objective, outcomes: list[str]) -> np.ndarray: +def extract_objective_weights(objective: Objective, outcomes: list[str]) -> npt.NDArray: """Extract a weights for objectives. Weights are for a maximization problem. @@ -475,18 +471,12 @@ def extract_outcome_constraints( def validate_and_apply_final_transform( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_weights: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - outcome_constraints: tuple[np.ndarray, np.ndarray] | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - linear_constraints: tuple[np.ndarray, np.ndarray] | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - pending_observations: list[np.ndarray] | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_thresholds: np.ndarray | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - final_transform: Callable[[np.ndarray], Tensor] = torch.tensor, + objective_weights: npt.NDArray, + outcome_constraints: tuple[npt.NDArray, npt.NDArray] | None, + linear_constraints: tuple[npt.NDArray, npt.NDArray] | None, + pending_observations: list[npt.NDArray] | None, + objective_thresholds: npt.NDArray | None = None, + final_transform: Callable[[npt.NDArray], Tensor] = torch.tensor, ) -> tuple[ Tensor, tuple[Tensor, Tensor] | None, @@ -558,8 +548,7 @@ def pending_observations_as_array_list( pending_observations: dict[str, list[ObservationFeatures]], outcome_names: list[str], param_names: list[str], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> list[np.ndarray] | None: +) -> list[npt.NDArray] | None: """Re-format pending observations. Args: @@ -591,8 +580,7 @@ def pending_observations_as_array_list( def parse_observation_features( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - X: np.ndarray, + X: npt.NDArray, param_names: list[str], candidate_metadata: list[TCandidateMetadata] | None = None, ) -> list[ObservationFeatures]: @@ -625,8 +613,7 @@ def parse_observation_features( def transform_callback( param_names: list[str], transforms: MutableMapping[str, Transform], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> Callable[[np.ndarray], np.ndarray]: +) -> Callable[[npt.NDArray], npt.NDArray]: """A closure for performing the `round trip` transformations. The function round points by de-transforming points back into @@ -643,8 +630,7 @@ def transform_callback( a function with for performing the roundtrip transform. """ - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _roundtrip_transform(x: np.ndarray) -> np.ndarray: + def _roundtrip_transform(x: npt.NDArray) -> npt.NDArray: """Inner function for performing aforementioned functionality. Args: @@ -1168,10 +1154,8 @@ def observed_hypervolume( def array_to_observation_data( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - f: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - cov: np.ndarray, + f: npt.NDArray, + cov: npt.NDArray, outcomes: list[str], ) -> list[ObservationData]: """Convert arrays of model predictions to a list of ObservationData. @@ -1198,8 +1182,7 @@ def array_to_observation_data( def observation_data_to_array( outcomes: list[str], observation_data: list[ObservationData], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[npt.NDArray, npt.NDArray]: """Convert a list of Observation data to arrays. Any missing mean or covariance values will be returned as NaNs. @@ -1235,19 +1218,15 @@ def observation_data_to_array( def observation_features_to_array( parameters: list[str], obsf: list[ObservationFeatures], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: +) -> npt.NDArray: """Convert a list of Observation features to arrays.""" return np.array([[of.parameters[p] for p in parameters] for of in obsf]) def feasible_hypervolume( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. optimization_config: MultiObjectiveOptimizationConfig, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - values: dict[str, np.ndarray], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: + values: dict[str, npt.NDArray], +) -> npt.NDArray: """Compute the feasible hypervolume each iteration. Args: @@ -1298,8 +1277,7 @@ def feasible_hypervolume( def _array_to_tensor( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - array: np.ndarray | list[float], + array: npt.NDArray | list[float], modelbridge: modelbridge_module.base.ModelBridge | None = None, ) -> Tensor: if modelbridge and hasattr(modelbridge, "_array_to_tensor"): diff --git a/ax/modelbridge/random.py b/ax/modelbridge/random.py index 412196cda73..da91c818f6f 100644 --- a/ax/modelbridge/random.py +++ b/ax/modelbridge/random.py @@ -26,8 +26,6 @@ FIT_MODEL_ERROR = "Model must be fit before {action}." -# pyre-fixme[13]: Attribute `model` is never initialized. -# pyre-fixme[13]: Attribute `parameters` is never initialized. class RandomModelBridge(ModelBridge): """A model bridge for using purely random 'models'. Data and optimization configs are not required. diff --git a/ax/modelbridge/tests/test_cross_validation.py b/ax/modelbridge/tests/test_cross_validation.py index 25a9ed47642..c5c9182359e 100644 --- a/ax/modelbridge/tests/test_cross_validation.py +++ b/ax/modelbridge/tests/test_cross_validation.py @@ -145,8 +145,6 @@ def test_CrossValidate(self) -> None: self.assertEqual(len(set(train[i]).intersection(test[i])), 0) self.assertEqual(len(train[i]) + len(test[i]), 4) # Test all points used as test points - # pyre-fixme[6]: For 1st param expected `Collection[ndarray]` but got - # `List[List[typing.Any]]`. all_test = np.hstack(test) self.assertTrue( np.array_equal(sorted(all_test), np.array([2.0, 2.0, 3.0, 4.0])) @@ -167,8 +165,6 @@ def test_CrossValidate(self) -> None: self.assertEqual(len(set(train[i]).intersection(test[i])), 0) self.assertEqual(len(train[i]) + len(test[i]), 4) # Test all points used as test points - # pyre-fixme[6]: For 1st param expected `Collection[ndarray]` but got - # `List[List[typing.Any]]`. all_test = np.hstack(test) self.assertTrue( np.array_equal(sorted(all_test), np.array([2.0, 2.0, 3.0, 4.0])) @@ -192,8 +188,6 @@ def test_CrossValidate(self) -> None: self.assertEqual(len(set(train[i]).intersection(test[i])), 0) self.assertEqual(len(train[i]) + len(test[i]), 4) # Test all points used as test points - # pyre-fixme[6]: For 1st param expected `Collection[ndarray]` but got - # `List[List[typing.Any]]`. all_test = np.hstack(test) self.assertTrue( np.array_equal(sorted(all_test), np.array([2.0, 2.0, 3.0, 4.0])) @@ -215,8 +209,6 @@ def test_selector(obs: Observation) -> bool: z = ma.cross_validate.mock_calls[5:] self.assertEqual(len(z), 2) all_test = np.hstack( - # pyre-fixme[6]: For 1st param expected `Collection[ndarray]` but got - # `List[List[typing.Any]]`. [[obsf.parameters["x"] for obsf in r[2]["cv_test_points"]] for r in z] ) self.assertTrue(np.array_equal(sorted(all_test), np.array([2.0, 2.0, 3.0]))) diff --git a/ax/modelbridge/tests/test_modelbridge_utils.py b/ax/modelbridge/tests/test_modelbridge_utils.py index 9d39cf319da..b95e8a9dad0 100644 --- a/ax/modelbridge/tests/test_modelbridge_utils.py +++ b/ax/modelbridge/tests/test_modelbridge_utils.py @@ -9,6 +9,7 @@ from dataclasses import dataclass import numpy as np +import numpy.typing as npt import torch from ax.core.metric import Metric from ax.core.objective import MultiObjective @@ -42,7 +43,7 @@ def test__array_to_tensor(self) -> None: @dataclass class MockModelbridge(ModelBridge): - def _array_to_tensor(self, array: np.ndarray | list[float]): + def _array_to_tensor(self, array: npt.NDArray | list[float]): return _array_to_tensor(array=array) mock_modelbridge = MockModelbridge() diff --git a/ax/modelbridge/tests/test_torch_modelbridge.py b/ax/modelbridge/tests/test_torch_modelbridge.py index a282ddddcad..c7b5655efd7 100644 --- a/ax/modelbridge/tests/test_torch_modelbridge.py +++ b/ax/modelbridge/tests/test_torch_modelbridge.py @@ -744,7 +744,7 @@ def test_convert_observations(self) -> None: feature_names=feature_names, bounds=[(0.0, 5.0)] * 3, ordinal_features=[2], - discrete_choices={2: list(range(0, 11))}, # pyre-ignore + discrete_choices={2: list(range(0, 11))}, task_features=[2] if use_task else [], target_values={2: 0} if use_task else {}, # pyre-ignore ) @@ -827,7 +827,7 @@ def test_convert_contextual_observations(self) -> None: feature_names=feature_names, bounds=[(0.0, 5.0)] * 3, ordinal_features=[2], - discrete_choices={2: list(range(0, 11))}, # pyre-ignore + discrete_choices={2: list(range(0, 11))}, task_features=[], target_values={}, ), diff --git a/ax/modelbridge/torch.py b/ax/modelbridge/torch.py index 79112e88c80..6b1099874e1 100644 --- a/ax/modelbridge/torch.py +++ b/ax/modelbridge/torch.py @@ -15,6 +15,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import torch from ax.core.arm import Arm from ax.core.data import Data @@ -76,7 +77,6 @@ FIT_MODEL_ERROR = "Model must be fit before {action}." -# pyre-fixme [13]: Attributes are never initialized. class TorchModelBridge(ModelBridge): """A model bridge for using torch-based models. @@ -283,22 +283,18 @@ def model_best_point( return best_arm, best_point_predictions def _array_callable_to_tensor_callable( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. self, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - array_func: Callable[[np.ndarray], np.ndarray], + array_func: Callable[[npt.NDArray], npt.NDArray], ) -> Callable[[Tensor], Tensor]: tensor_func: Callable[[Tensor], Tensor] = lambda x: ( self._array_to_tensor(array_func(x.detach().cpu().clone().numpy())) ) return tensor_func - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _array_list_to_tensors(self, arrays: list[np.ndarray]) -> list[Tensor]: + def _array_list_to_tensors(self, arrays: list[npt.NDArray]) -> list[Tensor]: return [self._array_to_tensor(x) for x in arrays] - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _array_to_tensor(self, array: np.ndarray | list[float]) -> Tensor: + def _array_to_tensor(self, array: npt.NDArray | list[float]) -> Tensor: return torch.as_tensor(array, dtype=self.dtype, device=self.device) def _convert_observations( @@ -766,10 +762,8 @@ def _predict( return array_to_observation_data(f=f, cov=cov, outcomes=self.outcomes) def _array_to_observation_features( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. self, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - X: np.ndarray, + X: npt.NDArray, candidate_metadata: list[TCandidateMetadata] | None, ) -> list[ObservationFeatures]: return parse_observation_features( diff --git a/ax/modelbridge/transforms/choice_encode.py b/ax/modelbridge/transforms/choice_encode.py index a507aadd4e2..bd95a177a4f 100644 --- a/ax/modelbridge/transforms/choice_encode.py +++ b/ax/modelbridge/transforms/choice_encode.py @@ -9,6 +9,7 @@ from typing import Any, Optional, TYPE_CHECKING import numpy as np +import numpy.typing as npt from ax.core.observation import Observation, ObservationFeatures from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter from ax.core.search_space import SearchSpace @@ -225,8 +226,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def transform_choice_values(p: ChoiceParameter) -> tuple[np.ndarray, ParameterType]: +def transform_choice_values(p: ChoiceParameter) -> tuple[npt.NDArray, ParameterType]: """Transforms the choice values and returns the new parameter type. If the choices were numeric (int or float) and ordered, then they're cast diff --git a/ax/modelbridge/transforms/ivw.py b/ax/modelbridge/transforms/ivw.py index 0edaa1def96..83911c51eb5 100644 --- a/ax/modelbridge/transforms/ivw.py +++ b/ax/modelbridge/transforms/ivw.py @@ -9,6 +9,7 @@ from logging import Logger import numpy as np +import numpy.typing as npt from ax.core.observation import ObservationData from ax.modelbridge.transforms.base import Transform from ax.utils.common.logger import get_logger @@ -48,8 +49,7 @@ def ivw_metric_merge( # weights is a map from metric name to a vector of the weights for each # measurement of that metric. indicies gives the corresponding index in # obsd.means for each measurement. - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - weights: dict[str, np.ndarray] = {} + weights: dict[str, npt.NDArray] = {} indicies: dict[str, list[int]] = {} for metric_name in set(obsd.metric_names): indcs = [i for i, mn in enumerate(obsd.metric_names) if mn == metric_name] @@ -64,7 +64,6 @@ def ivw_metric_merge( # Weight is inverse of variance, normalized # Expected `np.ndarray` for 3rd anonymous parameter to call # `dict.__setitem__` but got `float`. - # pyre-fixme[6]: weights[metric_name] = 1.0 / sigma2s weights[metric_name] /= np.sum(weights[metric_name]) else: @@ -99,8 +98,7 @@ def ivw_metric_merge( def _check_conflicting_means( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - means_noiseless: np.ndarray, + means_noiseless: npt.NDArray, metric_name: str, conflicting_noiseless: str, ) -> None: diff --git a/ax/modelbridge/transforms/log_y.py b/ax/modelbridge/transforms/log_y.py index e6a23672e62..e79851c2c23 100644 --- a/ax/modelbridge/transforms/log_y.py +++ b/ax/modelbridge/transforms/log_y.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING import numpy as np +import numpy.typing as npt from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.optimization_config import OptimizationConfig from ax.core.outcome_constraint import OutcomeConstraint @@ -102,8 +103,9 @@ def transform_optimization_config( def _tf_obs_data( self, observation_data: list[ObservationData], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - transform: Callable[[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]], + transform: Callable[ + [npt.NDArray, npt.NDArray], tuple[npt.NDArray, npt.NDArray] + ], ) -> list[ObservationData]: for obsd in observation_data: cov = obsd.covariance @@ -158,15 +160,11 @@ def untransform_outcome_constraints( def match_ci_width( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - mean: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - variance: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - transform: Callable[[np.ndarray], np.ndarray], + mean: npt.NDArray, + variance: npt.NDArray, + transform: Callable[[npt.NDArray], npt.NDArray], level: float = 0.95, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: +) -> npt.NDArray: fac = norm.ppf(1 - (1 - level) / 2) d = fac * np.sqrt(variance) width_asym = transform(mean + d) - transform(mean - d) @@ -177,12 +175,9 @@ def match_ci_width( def lognorm_to_norm( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - mu_ln: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Cov_ln: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[np.ndarray, np.ndarray]: + mu_ln: npt.NDArray, + Cov_ln: npt.NDArray, +) -> tuple[npt.NDArray, npt.NDArray]: """Compute mean and covariance of a MVN from those of the associated log-MVN If `Y` is log-normal with mean mu_ln and covariance Cov_ln, then @@ -197,12 +192,9 @@ def lognorm_to_norm( def norm_to_lognorm( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - mu_n: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Cov_n: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[np.ndarray, np.ndarray]: + mu_n: npt.NDArray, + Cov_n: npt.NDArray, +) -> tuple[npt.NDArray, npt.NDArray]: """Compute mean and covariance of a log-MVN from its MVN sufficient statistics If `X ~ N(mu_n, Cov_n)` and `Y = exp(X)`, then `Y` is log-normal with diff --git a/ax/modelbridge/transforms/metrics_as_task.py b/ax/modelbridge/transforms/metrics_as_task.py index f3473366853..3e162d6cc99 100644 --- a/ax/modelbridge/transforms/metrics_as_task.py +++ b/ax/modelbridge/transforms/metrics_as_task.py @@ -81,7 +81,7 @@ def transform_observations( new_obs_feats = obs.features.clone(replace_parameters=params) new_obs_data = ObservationData( metric_names=target_metrics, - means=obs.data.means_dict[task_metric] # pyre-ignore + means=obs.data.means_dict[task_metric] * np.ones(len(target_metrics)), covariance=np.diag( obs.data.covariance_matrix[task_metric][task_metric] diff --git a/ax/modelbridge/transforms/relativize.py b/ax/modelbridge/transforms/relativize.py index 1c207c2b6bb..896c1a56998 100644 --- a/ax/modelbridge/transforms/relativize.py +++ b/ax/modelbridge/transforms/relativize.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING import numpy as np +import numpy.typing as npt from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, @@ -170,8 +171,7 @@ def untransform_observations( def _get_relative_data_from_obs( self, obs: Observation, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], ) -> ObservationData: idx = ( int(obs.features.trial_index) @@ -192,8 +192,7 @@ def _get_relative_data_from_obs( def _rel_op_on_observations( self, observations: list[Observation], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], ) -> list[Observation]: return [ Observation( @@ -208,8 +207,7 @@ def _get_relative_data( self, data: ObservationData, status_quo_data: ObservationData, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], ) -> ObservationData: r""" Relativize or unrelativize `data` based on `status_quo_data` based on `rel_op` @@ -256,10 +254,8 @@ def _get_rel_mean_sem( mean_c: float, sem_c: float, metric: str, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ) -> tuple[float | np.ndarray, float | np.ndarray]: + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], + ) -> tuple[float | npt.NDArray, float | npt.NDArray]: """Compute (un)relativized mean and sem for a single metric.""" # if the is the status quo if means_t == mean_c and sems_t == sem_c: diff --git a/ax/modelbridge/transforms/rounding.py b/ax/modelbridge/transforms/rounding.py index f1b58562286..9f19db0e469 100644 --- a/ax/modelbridge/transforms/rounding.py +++ b/ax/modelbridge/transforms/rounding.py @@ -11,6 +11,7 @@ from copy import copy import numpy as np +import numpy.typing as npt from ax.core.parameter_constraint import OrderConstraint from ax.core.search_space import SearchSpace from ax.core.types import TParameterization @@ -22,8 +23,7 @@ def randomized_round(x: float) -> int: return int(z + float(random.random() <= (x - z))) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def randomized_onehot_round(x: np.ndarray) -> np.ndarray: +def randomized_onehot_round(x: npt.NDArray) -> npt.NDArray: """Randomized rounding of x to a one-hot vector. x should be 0 <= x <= 1. If x includes negative values, they will be rounded to zero. @@ -42,8 +42,7 @@ def randomized_onehot_round(x: np.ndarray) -> np.ndarray: return z -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def strict_onehot_round(x: np.ndarray) -> np.ndarray: +def strict_onehot_round(x: npt.NDArray) -> npt.NDArray: """Round x to a one-hot vector by selecting the max element. Ties broken randomly.""" if len(x) == 1: diff --git a/ax/modelbridge/transforms/tests/test_log_y_transform.py b/ax/modelbridge/transforms/tests/test_log_y_transform.py index 2393d0173d0..961711471c0 100644 --- a/ax/modelbridge/transforms/tests/test_log_y_transform.py +++ b/ax/modelbridge/transforms/tests/test_log_y_transform.py @@ -256,12 +256,10 @@ def test_lognorm_to_norm(self) -> None: def test_norm_to_lognorm(self) -> None: mu_n = -0.5 * np.ones(3) Cov_n = np.eye(3) - # pyre-fixme[6]: For 1st param expected `ndarray` but got `float`. mu_ln, Cov_ln = norm_to_lognorm(mu_n, Cov_n) self.assertTrue(np.allclose(mu_ln, np.ones(3))) self.assertTrue(np.allclose(Cov_ln, (np.exp(1) - 1) * np.eye(3))) Cov_n2 = np.array([[1.0, 0.0, 0.5], [0.0, 1.0, 0.0], [0.5, 0.0, 1.0]]) - # pyre-fixme[6]: For 1st param expected `ndarray` but got `float`. mu_ln2, Cov_ln2 = norm_to_lognorm(mu_n, Cov_n2) Z = np.zeros_like(Cov_ln2) Z[0, 2] = np.sqrt(np.exp(1)) - 1 diff --git a/ax/modelbridge/transforms/tests/test_relativize_transform.py b/ax/modelbridge/transforms/tests/test_relativize_transform.py index 32a485ceed9..8de382ef3fa 100644 --- a/ax/modelbridge/transforms/tests/test_relativize_transform.py +++ b/ax/modelbridge/transforms/tests/test_relativize_transform.py @@ -9,6 +9,7 @@ from unittest.mock import Mock import numpy as np +import numpy.typing as npt from ax.core import BatchTrial from ax.core.observation import ( Observation, @@ -48,8 +49,7 @@ class RelativizeDataTest(TestCase): Relativize, RelativizeWithConstantControl, ] - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - cases: list[tuple[type[Transform], list[tuple[np.ndarray, np.ndarray]]]] = [ + cases: list[tuple[type[Transform], list[tuple[npt.NDArray, npt.NDArray]]]] = [ ( Relativize, [ @@ -204,8 +204,7 @@ def test_relativize_transform_observations(self) -> None: def _check_transform_observations( tf: Transform, observations: list[Observation], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - expected_mean_and_covar: list[tuple[np.ndarray, np.ndarray]], + expected_mean_and_covar: list[tuple[npt.NDArray, npt.NDArray]], ) -> None: results = tf.transform_observations(observations) for i, tsfm_obs in enumerate(results): diff --git a/ax/modelbridge/transforms/tests/test_transform_to_new_sq.py b/ax/modelbridge/transforms/tests/test_transform_to_new_sq.py index 5802179a5cd..c931adf0ba0 100644 --- a/ax/modelbridge/transforms/tests/test_transform_to_new_sq.py +++ b/ax/modelbridge/transforms/tests/test_transform_to_new_sq.py @@ -7,6 +7,7 @@ import numpy as np +import numpy.typing as npt from ax.core.batch_trial import BatchTrial from ax.modelbridge import ModelBridge from ax.modelbridge.transforms.base import Transform @@ -28,8 +29,7 @@ class TransformToNewSQTest(RelativizeDataTest): # [Type[TransformToNewSQ]]` is not a subtype of the # overridden attribute `List[Type[Transform]]` relativize_classes = [TransformToNewSQ] - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - cases: list[tuple[type[Transform], list[tuple[np.ndarray, np.ndarray]]]] = [ + cases: list[tuple[type[Transform], list[tuple[npt.NDArray, npt.NDArray]]]] = [ ( TransformToNewSQ, [ diff --git a/ax/modelbridge/transforms/transform_to_new_sq.py b/ax/modelbridge/transforms/transform_to_new_sq.py index 53bd6d462c3..b349102bf45 100644 --- a/ax/modelbridge/transforms/transform_to_new_sq.py +++ b/ax/modelbridge/transforms/transform_to_new_sq.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING import numpy as np +import numpy.typing as npt from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.optimization_config import OptimizationConfig from ax.core.outcome_constraint import OutcomeConstraint @@ -86,8 +87,7 @@ def untransform_outcome_constraints( def _get_relative_data_from_obs( self, obs: Observation, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], ) -> ObservationData: idx = ( int(obs.features.trial_index) @@ -105,8 +105,7 @@ def _get_relative_data_from_obs( def _rel_op_on_observations( self, observations: list[Observation], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], ) -> list[Observation]: rel_observations = super()._rel_op_on_observations( observations=observations, @@ -126,8 +125,7 @@ def _get_relative_data( self, data: ObservationData, status_quo_data: ObservationData, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], ) -> ObservationData: r""" Transform or untransform `data` based on `status_quo_data` based on `rel_op`. @@ -174,8 +172,7 @@ def _get_rel_mean_sem( mean_c: float, sem_c: float, metric: str, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rel_op: Callable[..., tuple[np.ndarray, np.ndarray]], + rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], ) -> tuple[float, float]: """Compute (un)transformed mean and sem for a single metric.""" target_status_quo_data = self.status_quo_data_by_trial[self.default_trial_idx] diff --git a/ax/modelbridge/transforms/winsorize.py b/ax/modelbridge/transforms/winsorize.py index cb832da17a4..31bb1e5f2cf 100644 --- a/ax/modelbridge/transforms/winsorize.py +++ b/ax/modelbridge/transforms/winsorize.py @@ -11,6 +11,7 @@ from typing import Optional, TYPE_CHECKING import numpy as np +import numpy.typing as npt from ax.core.objective import MultiObjective, ScalarizedObjective from ax.core.observation import Observation, ObservationData from ax.core.optimization_config import ( @@ -344,8 +345,7 @@ def _get_objective_threshold_from_moo_config( ] -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _get_tukey_cutoffs(Y: np.ndarray, lower: bool) -> float: +def _get_tukey_cutoffs(Y: npt.NDArray, lower: bool) -> float: """Compute winsorization cutoffs similarly to Tukey boxplots. See https://mathworld.wolfram.com/Box-and-WhiskerPlot.html for more details. diff --git a/ax/models/discrete/full_factorial.py b/ax/models/discrete/full_factorial.py index 99317dfe0ac..04a7710bf89 100644 --- a/ax/models/discrete/full_factorial.py +++ b/ax/models/discrete/full_factorial.py @@ -11,7 +11,7 @@ from functools import reduce from operator import mul -import numpy as np +import numpy.typing as npt from ax.core.types import TGenMetadata, TParamValue, TParamValueList from ax.models.discrete_base import DiscreteModel from ax.models.types import TConfig @@ -52,10 +52,8 @@ def gen( self, n: int, parameter_values: list[TParamValueList], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_weights: np.ndarray | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - outcome_constraints: tuple[np.ndarray, np.ndarray] | None = None, + objective_weights: npt.NDArray | None, + outcome_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, fixed_features: dict[int, TParamValue] | None = None, pending_observations: list[list[TParamValueList]] | None = None, model_gen_options: TConfig | None = None, diff --git a/ax/models/discrete/thompson.py b/ax/models/discrete/thompson.py index 65e3bba5a67..27e7e044737 100644 --- a/ax/models/discrete/thompson.py +++ b/ax/models/discrete/thompson.py @@ -10,6 +10,7 @@ import json import numpy as np +import numpy.typing as npt from ax.core.types import TGenMetadata, TParamValue, TParamValueList from ax.exceptions.constants import TS_MIN_WEIGHT_ERROR, TS_NO_FEASIBLE_ARMS_ERROR from ax.exceptions.model import ModelError @@ -76,10 +77,8 @@ def gen( self, n: int, parameter_values: list[TParamValueList], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_weights: np.ndarray | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - outcome_constraints: tuple[np.ndarray, np.ndarray] | None = None, + objective_weights: npt.NDArray | None, + outcome_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, fixed_features: dict[int, TParamValue] | None = None, pending_observations: list[list[TParamValueList]] | None = None, model_gen_options: TConfig | None = None, @@ -124,8 +123,7 @@ def gen( return top_arms, top_weights, {"arms_to_weights": list(zip(arms, weights))} @copy_doc(DiscreteModel.predict) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def predict(self, X: list[TParamValueList]) -> tuple[np.ndarray, np.ndarray]: + def predict(self, X: list[TParamValueList]) -> tuple[npt.NDArray, npt.NDArray]: n = len(X) # number of parameterizations at which to make predictions m = len(self.Ys) # number of outcomes f = np.zeros((n, m)) # array of outcome predictions @@ -144,10 +142,8 @@ def predict(self, X: list[TParamValueList]) -> tuple[np.ndarray, np.ndarray]: def _generate_weights( self, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_weights: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - outcome_constraints: tuple[np.ndarray, np.ndarray] | None = None, + objective_weights: npt.NDArray, + outcome_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, ) -> list[float]: samples, fraction_all_infeasible = self._produce_samples( num_samples=self.num_samples, @@ -178,8 +174,7 @@ def _generate_weights( weights = winner_counts / winner_counts.sum() return weights.tolist() - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _generate_samples_per_metric(self, num_samples: int) -> np.ndarray: + def _generate_samples_per_metric(self, num_samples: int) -> npt.NDArray: k = len(self.X) samples_per_metric = np.zeros( (k, num_samples, len(self.Ys)) @@ -196,12 +191,9 @@ def _generate_samples_per_metric(self, num_samples: int) -> np.ndarray: def _produce_samples( self, num_samples: int, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_weights: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - outcome_constraints: tuple[np.ndarray, np.ndarray] | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ) -> tuple[np.ndarray, float]: + objective_weights: npt.NDArray, + outcome_constraints: tuple[npt.NDArray, npt.NDArray] | None, + ) -> tuple[npt.NDArray, float]: k = len(self.X) samples_per_metric = self._generate_samples_per_metric(num_samples=num_samples) diff --git a/ax/models/discrete_base.py b/ax/models/discrete_base.py index 8a4dd8a95e6..c6211ace7b6 100644 --- a/ax/models/discrete_base.py +++ b/ax/models/discrete_base.py @@ -7,7 +7,7 @@ # pyre-strict -import numpy as np +import numpy.typing as npt from ax.core.types import TGenMetadata, TParamValue, TParamValueList from ax.models.base import Model from ax.models.types import TConfig @@ -42,8 +42,7 @@ def fit( """ pass - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def predict(self, X: list[TParamValueList]) -> tuple[np.ndarray, np.ndarray]: + def predict(self, X: list[TParamValueList]) -> tuple[npt.NDArray, npt.NDArray]: """Predict Args: @@ -62,10 +61,8 @@ def gen( self, n: int, parameter_values: list[TParamValueList], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_weights: np.ndarray | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - outcome_constraints: tuple[np.ndarray, np.ndarray] | None = None, + objective_weights: npt.NDArray | None, + outcome_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, fixed_features: dict[int, TParamValue] | None = None, pending_observations: list[list[TParamValueList]] | None = None, model_gen_options: TConfig | None = None, @@ -105,8 +102,7 @@ def cross_validate( Yvars_train: list[list[float]], X_test: list[TParamValueList], use_posterior_predictive: bool = False, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ) -> tuple[np.ndarray, np.ndarray]: + ) -> tuple[npt.NDArray, npt.NDArray]: """Do cross validation with the given training and test sets. Training set is given in the same format as to fit. Test set is given @@ -137,10 +133,8 @@ def best_point( self, n: int, parameter_values: list[TParamValueList], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - objective_weights: np.ndarray | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - outcome_constraints: tuple[np.ndarray, np.ndarray] | None = None, + objective_weights: npt.NDArray | None, + outcome_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, fixed_features: dict[int, TParamValue] | None = None, pending_observations: list[list[TParamValueList]] | None = None, model_gen_options: TConfig | None = None, diff --git a/ax/models/model_utils.py b/ax/models/model_utils.py index e651d466792..bc1e68664b1 100644 --- a/ax/models/model_utils.py +++ b/ax/models/model_utils.py @@ -14,6 +14,7 @@ from typing import Protocol, Sequence, Union import numpy as np +import numpy.typing as npt import torch from ax.core.search_space import SearchSpaceDigest from ax.exceptions.core import SearchSpaceExhausted, UnsupportedError @@ -50,25 +51,19 @@ def predict(self, X: Tensor) -> tuple[Tensor, Tensor]: def rejection_sample( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. gen_unconstrained: Callable[ - [int, int, np.ndarray, dict[int, float] | None], np.ndarray + [int, int, npt.NDArray, dict[int, float] | None], npt.NDArray ], n: int, d: int, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - tunable_feature_indices: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - linear_constraints: tuple[np.ndarray, np.ndarray] | None = None, + tunable_feature_indices: npt.NDArray, + linear_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, deduplicate: bool = False, max_draws: int | None = None, fixed_features: dict[int, float] | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rounding_func: Callable[[np.ndarray], np.ndarray] | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - existing_points: np.ndarray | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[np.ndarray, int]: + rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None, + existing_points: npt.NDArray | None = None, +) -> tuple[npt.NDArray, int]: """Rejection sample in parameter space. Parameter space is typically [0, 1] for all tunable parameters. @@ -164,8 +159,7 @@ def rejection_sample( return (points, attempted_draws) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def check_duplicate(point: np.ndarray, points: np.ndarray) -> bool: +def check_duplicate(point: npt.NDArray, points: npt.NDArray) -> bool: """Check if a point exists in another array. Args: @@ -182,14 +176,11 @@ def check_duplicate(point: np.ndarray, points: np.ndarray) -> bool: def add_fixed_features( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - tunable_points: np.ndarray, + tunable_points: npt.NDArray, d: int, fixed_features: dict[int, float] | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - tunable_feature_indices: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: + tunable_feature_indices: npt.NDArray, +) -> npt.NDArray: """Add fixed features to points in tunable space. Args: @@ -213,12 +204,9 @@ def add_fixed_features( def check_param_constraints( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - linear_constraints: tuple[np.ndarray, np.ndarray], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - point: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[bool, np.ndarray]: + linear_constraints: tuple[npt.NDArray, npt.NDArray], + point: npt.NDArray, +) -> tuple[bool, npt.NDArray]: """Check if a point satisfies parameter constraints. Args: @@ -245,8 +233,7 @@ def check_param_constraints( def tunable_feature_indices( bounds: list[tuple[float, float]], fixed_features: dict[int, float] | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: +) -> npt.NDArray: """Get the feature indices of tunable features. Args: @@ -263,10 +250,8 @@ def tunable_feature_indices( def validate_bounds( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. bounds: list[tuple[float, float]], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - fixed_feature_indices: np.ndarray, + fixed_feature_indices: npt.NDArray, ) -> None: """Ensure the requested space is [0,1]^d. @@ -364,8 +349,7 @@ def best_observed_point( def best_in_sample_point( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Xs: list[torch.Tensor] | list[np.ndarray], + Xs: list[torch.Tensor] | list[npt.NDArray], model: TorchModelLike, bounds: list[tuple[float, float]], objective_weights: Tensoray | None, @@ -492,8 +476,9 @@ def best_in_sample_point( return X_obs[i, :], utility[i] -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def as_array(x: Tensoray | tuple[Tensoray, ...]) -> np.ndarray | tuple[np.ndarray, ...]: +def as_array( + x: Tensoray | tuple[Tensoray, ...], +) -> npt.NDArray | tuple[npt.NDArray, ...]: """Convert every item in a tuple of tensors/arrays into an array. Args: @@ -513,8 +498,7 @@ def as_array(x: Tensoray | tuple[Tensoray, ...]) -> np.ndarray | tuple[np.ndarra def get_observed( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Xs: list[torch.Tensor] | list[np.ndarray], + Xs: list[torch.Tensor] | list[npt.NDArray], objective_weights: Tensoray, outcome_constraints: tuple[Tensoray, Tensoray] | None = None, # pyre-fixme[7]: Expected `Union[ndarray[typing.Any, typing.Any], Tensor]` but got @@ -553,8 +537,6 @@ def get_observed( # pyre-fixme[6]: For 2nd param expected `Union[None, Dict[str, Tuple[typing.A... return np.array(list(X_obs_set), dtype=Xs[0].dtype) # (n x d) if isinstance(Xs[0], torch.Tensor): - # pyre-fixme[7]: Expected `Union[np.ndarray, torch.Tensor]` but got implicit - # return value of `None`. # pyre-fixme[6]: For 3rd param expected `Optional[_C.dtype]` but got # `Union[np.dtype, _C.dtype]`. # pyre-fixme[16]: Item `ndarray` of `Union[ndarray[typing.Any, typing.Any], diff --git a/ax/models/random/base.py b/ax/models/random/base.py index 8fb0800ae2a..1a6f5ff7e25 100644 --- a/ax/models/random/base.py +++ b/ax/models/random/base.py @@ -11,6 +11,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import torch from ax.exceptions.core import SearchSpaceExhausted from ax.models.base import Model @@ -65,8 +66,7 @@ def __init__( deduplicate: bool = True, seed: int | None = None, init_position: int = 0, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - generated_points: np.ndarray | None = None, + generated_points: npt.NDArray | None = None, fallback_to_sample_polytope: bool = False, ) -> None: super().__init__() @@ -86,14 +86,11 @@ def gen( self, n: int, bounds: list[tuple[float, float]], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - linear_constraints: tuple[np.ndarray, np.ndarray] | None = None, + linear_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, fixed_features: dict[int, float] | None = None, model_gen_options: TConfig | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rounding_func: Callable[[np.ndarray], np.ndarray] | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ) -> tuple[np.ndarray, np.ndarray]: + rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None, + ) -> tuple[npt.NDArray, npt.NDArray]: """Generate new candidates. Args: @@ -204,11 +201,9 @@ def _gen_unconstrained( self, n: int, d: int, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - tunable_feature_indices: np.ndarray, + tunable_feature_indices: npt.NDArray, fixed_features: dict[int, float] | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ) -> np.ndarray: + ) -> npt.NDArray: """Generate n points, from an unconstrained parameter space, using _gen_samples. Args: @@ -231,8 +226,7 @@ def _gen_unconstrained( ) return points - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _gen_samples(self, n: int, tunable_d: int) -> np.ndarray: + def _gen_samples(self, n: int, tunable_d: int) -> npt.NDArray: """Generate n samples on [0, 1]^d. Args: @@ -245,10 +239,8 @@ def _gen_samples(self, n: int, tunable_d: int) -> np.ndarray: raise NotImplementedError("Base RandomModel can't generate samples.") def _convert_inequality_constraints( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. self, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - linear_constraints: tuple[np.ndarray, np.ndarray] | None, + linear_constraints: tuple[npt.NDArray, npt.NDArray] | None, ) -> tuple[Tensor, Tensor] | None: """Helper method to convert inequality constraints used by the rejection sampler to the format required for the polytope sampler. diff --git a/ax/models/random/sobol.py b/ax/models/random/sobol.py index a2eb6aa6b16..28d7fbddfd3 100644 --- a/ax/models/random/sobol.py +++ b/ax/models/random/sobol.py @@ -9,6 +9,7 @@ from collections.abc import Callable import numpy as np +import numpy.typing as npt import torch from ax.models.model_utils import tunable_feature_indices from ax.models.random.base import RandomModel @@ -35,8 +36,7 @@ def __init__( seed: int | None = None, init_position: int = 0, scramble: bool = True, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - generated_points: np.ndarray | None = None, + generated_points: npt.NDArray | None = None, fallback_to_sample_polytope: bool = False, ) -> None: super().__init__( @@ -76,14 +76,11 @@ def gen( self, n: int, bounds: list[tuple[float, float]], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - linear_constraints: tuple[np.ndarray, np.ndarray] | None = None, + linear_constraints: tuple[npt.NDArray, npt.NDArray] | None = None, fixed_features: dict[int, float] | None = None, model_gen_options: TConfig | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - rounding_func: Callable[[np.ndarray], np.ndarray] | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ) -> tuple[np.ndarray, np.ndarray]: + rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None, + ) -> tuple[npt.NDArray, npt.NDArray]: """Generate new candidates. Args: @@ -121,8 +118,7 @@ def gen( self.init_position = not_none(self.engine).num_generated return points, weights - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _gen_samples(self, n: int, tunable_d: int) -> np.ndarray: + def _gen_samples(self, n: int, tunable_d: int) -> npt.NDArray: """Generate n samples. Args: diff --git a/ax/models/random/uniform.py b/ax/models/random/uniform.py index c8c8e72b52a..ceb4631527d 100644 --- a/ax/models/random/uniform.py +++ b/ax/models/random/uniform.py @@ -8,6 +8,7 @@ import numpy as np +import numpy.typing as npt from ax.models.random.base import RandomModel @@ -25,8 +26,7 @@ def __init__( deduplicate: bool = True, seed: int | None = None, init_position: int = 0, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - generated_points: np.ndarray | None = None, + generated_points: npt.NDArray | None = None, fallback_to_sample_polytope: bool = False, ) -> None: super().__init__( @@ -41,8 +41,7 @@ def __init__( # Fast-forward the random state by generating & discarding samples. self._rs.uniform(size=(self.init_position)) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _gen_samples(self, n: int, tunable_d: int) -> np.ndarray: + def _gen_samples(self, n: int, tunable_d: int) -> npt.NDArray: """Generate samples from the scipy uniform distribution. Args: diff --git a/ax/models/torch/botorch.py b/ax/models/torch/botorch.py index dd0bbafe13e..b737bca145d 100644 --- a/ax/models/torch/botorch.py +++ b/ax/models/torch/botorch.py @@ -14,7 +14,7 @@ from logging import Logger from typing import Any, Optional -import numpy as np +import numpy.typing as npt import torch from ax.core.search_space import SearchSpaceDigest from ax.core.types import TCandidateMetadata @@ -496,8 +496,7 @@ def cross_validate( # pyre-ignore [14]: `search_space_digest` arg not needed he model=model, X=X_test, use_posterior_predictive=use_posterior_predictive ) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def feature_importances(self) -> np.ndarray: + def feature_importances(self) -> npt.NDArray: return get_feature_importances_from_botorch_model(model=self._model) @property @@ -544,8 +543,7 @@ def botorch_rounding_func(X: Tensor) -> Tensor: def get_feature_importances_from_botorch_model( model: Model | ModuleList | None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: +) -> npt.NDArray: """Get feature importances from a list of BoTorch models. Args: diff --git a/ax/models/torch/botorch_modular/model.py b/ax/models/torch/botorch_modular/model.py index c48165377a3..a5e8f84aedc 100644 --- a/ax/models/torch/botorch_modular/model.py +++ b/ax/models/torch/botorch_modular/model.py @@ -13,7 +13,7 @@ from dataclasses import dataclass, field from typing import Any -import numpy as np +import numpy.typing as npt import torch from ax.core.search_space import SearchSpaceDigest from ax.core.types import TCandidateMetadata, TGenMetadata @@ -464,8 +464,7 @@ def _instantiate_acquisition( options=acq_options, ) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def feature_importances(self) -> np.ndarray: + def feature_importances(self) -> npt.NDArray: """Compute feature importances from the model. This assumes that we can get model lengthscales from either diff --git a/ax/models/torch/randomforest.py b/ax/models/torch/randomforest.py index 0bd9be205a3..4ac405bffcb 100644 --- a/ax/models/torch/randomforest.py +++ b/ax/models/torch/randomforest.py @@ -9,6 +9,7 @@ from __future__ import annotations import numpy as np +import numpy.typing as npt import torch from ax.core.search_space import SearchSpaceDigest from ax.core.types import TCandidateMetadata @@ -87,12 +88,9 @@ def cross_validate( # pyre-ignore [14]: not using metric_names or ssd def _get_rf( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - X: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Y: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Yvar: np.ndarray, + X: npt.NDArray, + Y: npt.NDArray, + Yvar: npt.NDArray, num_trees: int, max_features: str | None, ) -> RandomForestRegressor: diff --git a/ax/models/torch/tests/test_acquisition.py b/ax/models/torch/tests/test_acquisition.py index df13d29a91b..033484e3db2 100644 --- a/ax/models/torch/tests/test_acquisition.py +++ b/ax/models/torch/tests/test_acquisition.py @@ -444,8 +444,6 @@ def test_optimize_discrete(self) -> None: feature_names=["a", "b", "c"], bounds=[(0, 4) for _ in range(3)], categorical_features=[0, 1, 2], - # pyre-fixme[6]: For 4th param expected `Dict[int, List[Union[float, - # int]]]` but got `Dict[int, List[int]]`. discrete_choices={k: [0, 1, 2, 3, 4] for k in range(3)}, ) with mock.patch( diff --git a/ax/models/torch/tests/test_sebo.py b/ax/models/torch/tests/test_sebo.py index 18ffa4662f7..498ac0967ec 100644 --- a/ax/models/torch/tests/test_sebo.py +++ b/ax/models/torch/tests/test_sebo.py @@ -164,9 +164,7 @@ def test_init(self) -> None: # Check transformed objective threshold self.assertTrue( torch.equal( - # pyre-fixme[6]: For 2nd argument expected `Tensor` but got `int`. acquisition1.acqf.ref_point[-1], - # pyre-fixme[6]: For 2nd argument expected `Tensor` but got `int`. -self.objective_thresholds_sebo[-1], ) ) diff --git a/ax/models/torch/utils.py b/ax/models/torch/utils.py index bde273facd1..3009ea460d4 100644 --- a/ax/models/torch/utils.py +++ b/ax/models/torch/utils.py @@ -12,6 +12,7 @@ from typing import Any, cast import numpy as np +import numpy.typing as npt import torch from ax.exceptions.core import UnsupportedError from ax.exceptions.model import ModelError @@ -364,12 +365,10 @@ def _to_inequality_constraints( def tensor_callable_to_array_callable( tensor_func: Callable[[Tensor], Tensor], device: torch.device, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> Callable[[np.ndarray], np.ndarray]: +) -> Callable[[npt.NDArray], npt.NDArray]: """transfer a tensor callable to an array callable""" - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def array_func(x: np.ndarray) -> np.ndarray: + def array_func(x: npt.NDArray) -> npt.NDArray: return tensor_func(torch.from_numpy(x).to(device)).detach().cpu().numpy() return array_func diff --git a/ax/plot/contour.py b/ax/plot/contour.py index 30e767ff5cf..9970b635a3c 100644 --- a/ax/plot/contour.py +++ b/ax/plot/contour.py @@ -11,6 +11,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import plotly.graph_objs as go from ax.core.observation import ObservationFeatures from ax.modelbridge.base import ModelBridge @@ -401,8 +402,7 @@ def interact_contour_plotly( param_names = [parameter.name for parameter in range_parameters] is_log_dict: dict[str, bool] = {} - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - grid_dict: dict[str, np.ndarray] = {} + grid_dict: dict[str, npt.NDArray] = {} for parameter in range_parameters: is_log_dict[parameter.name] = parameter.log_scale grid_dict[parameter.name] = get_grid_for_parameter(parameter, density) @@ -413,14 +413,12 @@ def interact_contour_plotly( # pyre-fixme[9]: f_dict has type `Dict[str, Dict[str, np.ndarray]]`; used as # `Dict[str, Dict[str, typing.List[Variable[_T]]]]`. - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - f_dict: dict[str, dict[str, np.ndarray]] = { + f_dict: dict[str, dict[str, npt.NDArray]] = { param1: {param2: [] for param2 in param_names} for param1 in param_names } # pyre-fixme[9]: sd_dict has type `Dict[str, Dict[str, np.ndarray]]`; used as # `Dict[str, Dict[str, typing.List[Variable[_T]]]]`. - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sd_dict: dict[str, dict[str, np.ndarray]] = { + sd_dict: dict[str, dict[str, npt.NDArray]] = { param1: {param2: [] for param2 in param_names} for param1 in param_names } diff --git a/ax/plot/feature_importances.py b/ax/plot/feature_importances.py index 504d8f3cea8..97b5e91c607 100644 --- a/ax/plot/feature_importances.py +++ b/ax/plot/feature_importances.py @@ -10,6 +10,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import pandas as pd import plotly.graph_objs as go from ax.core.parameter import ChoiceParameter @@ -102,8 +103,7 @@ def plot_feature_importance_by_metric(model: ModelBridge) -> AxPlotConfig: def plot_feature_importance_by_feature_plotly( model: ModelBridge | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sensitivity_values: dict[str, dict[str, float | np.ndarray]] | None = None, + sensitivity_values: dict[str, dict[str, float | npt.NDArray]] | None = None, relative: bool = False, caption: str = "", importance_measure: str = "", @@ -280,8 +280,7 @@ def plot_feature_importance_by_feature_plotly( def plot_feature_importance_by_feature( model: ModelBridge | None = None, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sensitivity_values: dict[str, dict[str, float | np.ndarray]] | None = None, + sensitivity_values: dict[str, dict[str, float | npt.NDArray]] | None = None, relative: bool = False, caption: str = "", importance_measure: str = "", diff --git a/ax/plot/helper.py b/ax/plot/helper.py index a37f2da9792..bc7c3d3f8f6 100644 --- a/ax/plot/helper.py +++ b/ax/plot/helper.py @@ -14,6 +14,7 @@ from typing import Any, Optional, Union import numpy as np +import numpy.typing as npt from ax.core.generator_run import GeneratorRun from ax.core.observation import Observation, ObservationFeatures from ax.core.parameter import ChoiceParameter, FixedParameter, Parameter, RangeParameter @@ -460,8 +461,7 @@ def get_range_parameters( ) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def get_grid_for_parameter(parameter: RangeParameter, density: int) -> np.ndarray: +def get_grid_for_parameter(parameter: RangeParameter, density: int) -> npt.NDArray: """Get a grid of points along the range of the parameter. Will be a log-scale grid if parameter is log scale. diff --git a/ax/plot/pareto_frontier.py b/ax/plot/pareto_frontier.py index 02c2b63862e..cdf9551a520 100644 --- a/ax/plot/pareto_frontier.py +++ b/ax/plot/pareto_frontier.py @@ -10,6 +10,7 @@ from collections.abc import Iterable import numpy as np +import numpy.typing as npt import pandas as pd import plotly.graph_objs as go from ax.core.experiment import Experiment @@ -47,8 +48,7 @@ def _make_label( return f"{name}: {estimate}{perc} {ci}
" -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _filter_outliers(Y: np.ndarray, m: float = 2.0) -> np.ndarray: +def _filter_outliers(Y: npt.NDArray, m: float = 2.0) -> npt.NDArray: std_filter = abs(Y - np.median(Y, axis=0)) < m * np.std(Y, axis=0) return Y[np.all(abs(std_filter), axis=1)] @@ -81,10 +81,8 @@ def scatter_plot_with_hypervolume_trace_plotly(experiment: Experiment) -> go.Fig def scatter_plot_with_pareto_frontier_plotly( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Y: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Y_pareto: np.ndarray | None, + Y: npt.NDArray, + Y_pareto: npt.NDArray | None, metric_x: str | None, metric_y: str | None, reference_point: tuple[float, float] | None, @@ -241,10 +239,8 @@ def scatter_plot_with_pareto_frontier_plotly( def scatter_plot_with_pareto_frontier( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Y: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Y_pareto: np.ndarray, + Y: npt.NDArray, + Y_pareto: npt.NDArray, metric_x: str, metric_y: str, reference_point: tuple[float, float], diff --git a/ax/plot/pareto_utils.py b/ax/plot/pareto_utils.py index 1a3c12f75c6..2ae6a0afcf8 100644 --- a/ax/plot/pareto_utils.py +++ b/ax/plot/pareto_utils.py @@ -12,6 +12,7 @@ from typing import NamedTuple import numpy as np +import numpy.typing as npt import torch from ax.core.batch_trial import BatchTrial from ax.core.data import Data @@ -54,12 +55,10 @@ def _extract_observed_pareto_2d( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - Y: np.ndarray, + Y: npt.NDArray, reference_point: tuple[float, float] | None, minimize: bool | tuple[bool, bool] = True, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray: +) -> npt.NDArray: if Y.shape[1] != 2: raise NotImplementedError("Currently only the 2-dim case is handled.") # If `minimize` is a bool, apply to both dimensions @@ -553,8 +552,7 @@ def _validate_outcome_constraints( def _build_scalarized_optimization_config( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - weights: np.ndarray, + weights: npt.NDArray, primary_objective: Metric, secondary_objective: Metric, outcome_constraints: list[OutcomeConstraint] | None = None, diff --git a/ax/plot/slice.py b/ax/plot/slice.py index 05a94091ced..29b1c68ebcb 100644 --- a/ax/plot/slice.py +++ b/ax/plot/slice.py @@ -10,6 +10,7 @@ from typing import Any, Optional, Union import numpy as np +import numpy.typing as npt from ax.core.observation import ObservationFeatures from ax.modelbridge.base import ModelBridge from ax.plot.base import AxPlotConfig, AxPlotTypes, PlotData @@ -342,8 +343,7 @@ def interact_slice_plotly( plot_data_dict = {} raw_data_dict = {} - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sd_plt_dict: dict[str, dict[str, np.ndarray]] = {} + sd_plt_dict: dict[str, dict[str, npt.NDArray]] = {} cond_name_to_parameters_dict = {} is_log_dict: dict[str, bool] = {} diff --git a/ax/plot/trace.py b/ax/plot/trace.py index eb622a1dcb5..ed2858eceed 100644 --- a/ax/plot/trace.py +++ b/ax/plot/trace.py @@ -10,6 +10,7 @@ from typing import Any import numpy as np +import numpy.typing as npt import pandas as pd import plotly.graph_objs as go from ax.core.experiment import Experiment @@ -28,10 +29,8 @@ def map_data_single_trace_scatters( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - x: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y: np.ndarray, + x: npt.NDArray, + y: npt.NDArray, legend_label: str, xlabel: str = "Trial progression", ylabel: str = "Trial performance", @@ -99,10 +98,8 @@ def map_data_single_trace_scatters( def map_data_multiple_metrics_dropdown_plotly( title: str, metric_names: list[str], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - xs_by_metric: dict[str, list[np.ndarray]], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ys_by_metric: dict[str, list[np.ndarray]], + xs_by_metric: dict[str, list[npt.NDArray]], + ys_by_metric: dict[str, list[npt.NDArray]], legend_labels_by_metric: dict[str, list[str]], stopping_markers_by_metric: dict[str, list[bool]], xlabels_by_metric: dict[str, str], @@ -213,8 +210,7 @@ def map_data_multiple_metrics_dropdown_plotly( def mean_trace_scatter( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y: np.ndarray, + y: npt.NDArray, trace_color: tuple[int] = COLORS.STEELBLUE.value, legend_label: str = "mean", hover_labels: list[str] | None = None, @@ -247,8 +243,7 @@ def mean_trace_scatter( def sem_range_scatter( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y: np.ndarray, + y: npt.NDArray, trace_color: tuple[int] = COLORS.STEELBLUE.value, legend_label: str = "", ) -> tuple[go.Scatter, go.Scatter]: @@ -291,8 +286,7 @@ def sem_range_scatter( def mean_markers_scatter( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y: np.ndarray, + y: npt.NDArray, marker_color: tuple[int] = COLORS.LIGHT_PURPLE.value, legend_label: str = "", hover_labels: list[str] | None = None, @@ -353,8 +347,7 @@ def optimum_objective_scatter( def optimization_trace_single_method_plotly( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y: np.ndarray, + y: npt.NDArray, optimum: float | None = None, model_transitions: list[int] | None = None, title: str = "", @@ -457,8 +450,7 @@ def optimization_trace_single_method_plotly( def _autoset_axis_limits( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y: np.ndarray, + y: npt.NDArray, optimization_direction: str, force_include_value: float | None = None, ) -> list[float]: @@ -490,8 +482,7 @@ def _autoset_axis_limits( def optimization_trace_single_method( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y: np.ndarray, + y: npt.NDArray, optimum: float | None = None, model_transitions: list[int] | None = None, title: str = "", @@ -558,8 +549,7 @@ def optimization_trace_single_method( def optimization_trace_all_methods( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_dict: dict[str, np.ndarray], + y_dict: dict[str, npt.NDArray], optimum: float | None = None, title: str = "", ylabel: str = "", diff --git a/ax/service/tests/scheduler_test_utils.py b/ax/service/tests/scheduler_test_utils.py index e8f8b7352a6..f80f38ffc2b 100644 --- a/ax/service/tests/scheduler_test_utils.py +++ b/ax/service/tests/scheduler_test_utils.py @@ -959,7 +959,6 @@ def test_stop_at_MAX_SECONDS_BETWEEN_REPORTS(self) -> None: with patch.object( scheduler, "wait_for_completed_trials_and_report_results", return_value=None ) as mock_await_trials: - # pyre-fixme[6]: For 1st param expected `Optional[int]` but got `float`. scheduler.run_all_trials(timeout_hours=1 / 60 / 15) # 4 second timeout. # We should be calling `wait_for_completed_trials_and_report_results` # N = total runtime / `test_stop_at_MAX_SECONDS_BETWEEN_REPORTS` times. diff --git a/ax/service/tests/test_instantiation_utils.py b/ax/service/tests/test_instantiation_utils.py index 4cbf347897a..bdaddf20930 100644 --- a/ax/service/tests/test_instantiation_utils.py +++ b/ax/service/tests/test_instantiation_utils.py @@ -49,8 +49,6 @@ def test_constraint_from_str(self) -> None: with self.assertRaisesRegex(ValueError, "Bound for the constraint"): InstantiationBase.constraint_from_str( "x1 + x2 <= not_numerical_bound", - # pyre-fixme[6]: For 2nd param expected `Dict[str, Parameter]` but - # got `Dict[str, None]`. { "x1": RangeParameter( name="x1", @@ -103,8 +101,6 @@ def test_constraint_from_str(self) -> None: ) one_val_constraint = InstantiationBase.constraint_from_str( "x1 <= 0", - # pyre-fixme[6]: For 2nd param expected `Dict[str, Parameter]` but - # got `Dict[str, None]`. { "x1": RangeParameter( name="x1", parameter_type=ParameterType.FLOAT, lower=0.1, upper=2.0 @@ -118,8 +114,6 @@ def test_constraint_from_str(self) -> None: self.assertEqual(one_val_constraint.constraint_dict, {"x1": 1.0}) one_val_constraint = InstantiationBase.constraint_from_str( "-0.5*x1 >= -0.1", - # pyre-fixme[6]: For 2nd param expected `Dict[str, Parameter]` but - # got `Dict[str, None]`. { "x1": RangeParameter( name="x1", parameter_type=ParameterType.FLOAT, lower=0.1, upper=2.0 @@ -153,8 +147,6 @@ def test_constraint_from_str(self) -> None: with self.assertRaisesRegex(ValueError, "Multiplier should be float"): InstantiationBase.constraint_from_str( "x1 - e*x2 + x3 <= 3", - # pyre-fixme[6]: For 2nd param expected `Dict[str, Parameter]` but - # got `Dict[str, None]`. { "x1": RangeParameter( name="x1", @@ -179,8 +171,6 @@ def test_constraint_from_str(self) -> None: with self.assertRaisesRegex(ValueError, "A linear constraint should be"): InstantiationBase.constraint_from_str( "x1 - 2 *x2 + 3 *x3 <= 3", - # pyre-fixme[6]: For 2nd param expected `Dict[str, Parameter]` but - # got `Dict[str, None]`. { "x1": RangeParameter( name="x1", @@ -205,8 +195,6 @@ def test_constraint_from_str(self) -> None: with self.assertRaisesRegex(ValueError, "A linear constraint should be"): InstantiationBase.constraint_from_str( "x1 - 2* x2 + 3* x3 <= 3", - # pyre-fixme[6]: For 2nd param expected `Dict[str, Parameter]` but - # got `Dict[str, None]`. { "x1": RangeParameter( name="x1", @@ -231,8 +219,6 @@ def test_constraint_from_str(self) -> None: with self.assertRaisesRegex(ValueError, "A linear constraint should be"): InstantiationBase.constraint_from_str( "x1 - 2 * x2 + 3*x3 <= 3", - # pyre-fixme[6]: For 2nd param expected `Dict[str, Parameter]` but - # got `Dict[str, None]`. { "x1": RangeParameter( name="x1", diff --git a/ax/service/tests/test_managed_loop.py b/ax/service/tests/test_managed_loop.py index d20e34eaca0..716638c4157 100644 --- a/ax/service/tests/test_managed_loop.py +++ b/ax/service/tests/test_managed_loop.py @@ -9,6 +9,7 @@ from unittest.mock import Mock, patch import numpy as np +import numpy.typing as npt from ax.exceptions.core import UserInputError from ax.metrics.branin import branin from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy @@ -16,15 +17,13 @@ from ax.service.managed_loop import OptimizationLoop, optimize from ax.utils.common.testutils import TestCase from ax.utils.testing.mock import fast_botorch_optimize -from numpy import ndarray def _branin_evaluation_function( # pyre-fixme[2]: Parameter must be annotated. parameterization, weight=None, # pyre-fixme[2]: Parameter must be annotated. - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. -) -> dict[str, tuple[float | ndarray, float]]: +) -> dict[str, tuple[float | npt.NDArray, float]]: if any(param_name not in parameterization.keys() for param_name in ["x1", "x2"]): raise ValueError("Parametrization does not contain x1 or x2") x1, x2 = parameterization["x1"], parameterization["x2"] @@ -38,8 +37,7 @@ def _branin_evaluation_function_v2( # pyre-fixme[2]: Parameter must be annotated. parameterization, weight=None, # pyre-fixme[2]: Parameter must be annotated. - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. -) -> tuple[float | ndarray, float]: +) -> tuple[float | npt.NDArray, float]: if any(param_name not in parameterization.keys() for param_name in ["x1", "x2"]): raise ValueError("Parametrization does not contain x1 or x2") x1, x2 = parameterization["x1"], parameterization["x2"] @@ -50,8 +48,7 @@ def _branin_evaluation_function_with_unknown_sem( # pyre-fixme[2]: Parameter must be annotated. parameterization, weight=None, # pyre-fixme[2]: Parameter must be annotated. - # pyre-fixme[24]: Generic type `ndarray` expects 2 type parameters. -) -> tuple[float | ndarray, None]: +) -> tuple[float | npt.NDArray, None]: if any(param_name not in parameterization.keys() for param_name in ["x1", "x2"]): raise ValueError("Parametrization does not contain x1 or x2") x1, x2 = parameterization["x1"], parameterization["x2"] diff --git a/ax/service/utils/report_utils.py b/ax/service/utils/report_utils.py index c300fd918a1..5f3aa420e35 100644 --- a/ax/service/utils/report_utils.py +++ b/ax/service/utils/report_utils.py @@ -18,6 +18,7 @@ import gpytorch import numpy as np +import numpy.typing as npt import pandas as pd import plotly.graph_objects as go from ax.core.base_trial import TrialStatus @@ -148,9 +149,8 @@ def _get_objective_trace_plot( def _get_objective_v_param_plots( experiment: Experiment, model: ModelBridge, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. importance: None - | (dict[str, dict[str, np.ndarray]] | dict[str, dict[str, float]]) = None, + | (dict[str, dict[str, npt.NDArray]] | dict[str, dict[str, float]]) = None, # Chosen to take ~1min on local benchmarks. max_num_slice_plots: int = 200, # Chosen to take ~2min on local benchmarks. @@ -525,12 +525,10 @@ def get_standard_plots( def _transform_progression_to_walltime( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - progressions: np.ndarray, + progressions: npt.NDArray, exp_df: pd.DataFrame, trial_idx: int, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray | None: +) -> npt.NDArray | None: try: trial_df = exp_df[exp_df["trial_index"] == trial_idx] time_run_started = trial_df["time_run_started"].iloc[0] diff --git a/ax/utils/measurement/synthetic_functions.py b/ax/utils/measurement/synthetic_functions.py index 4d433ecb160..844dd689b50 100644 --- a/ax/utils/measurement/synthetic_functions.py +++ b/ax/utils/measurement/synthetic_functions.py @@ -10,6 +10,7 @@ from typing import TypeVar import numpy as np +import numpy.typing as npt import torch from ax.utils.common.docutils import copy_doc from ax.utils.common.typeutils import checked_cast, not_none @@ -38,12 +39,9 @@ def name(self) -> str: def __call__( self, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - *args: int | float | np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - **kwargs: int | float | np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - ) -> float | np.ndarray: + *args: int | float | npt.NDArray, + **kwargs: int | float | npt.NDArray, + ) -> float | npt.NDArray: """Simplified way to call the synthetic function and pass the argument numbers directly, e.g. `branin(2.0, 3.0)`. """ @@ -71,8 +69,7 @@ def __call__( x = float(x) return checked_cast(float, self.f(np.array(args))) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def f(self, X: np.ndarray) -> float | np.ndarray: + def f(self, X: npt.NDArray) -> float | npt.NDArray: """Synthetic function implementation. Args: @@ -148,8 +145,7 @@ def fmax(self) -> float: return self.informative_failure_on_none(self._fmax) @abstractmethod - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _f(self, X: np.ndarray) -> float: + def _f(self, X: npt.NDArray) -> float: """Implementation of the synthetic function. Must be implemented in subclass. Args: @@ -177,8 +173,7 @@ def name(self) -> str: return f"{self.__class__.__name__}_{self._botorch_function.__class__.__name__}" @override - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _f(self, X: np.ndarray) -> float: + def _f(self, X: npt.NDArray) -> float: # TODO: support batch evaluation return float(self._botorch_function(X=torch.from_numpy(X)).item()) @@ -198,10 +193,8 @@ class Hartmann6(SyntheticFunction): _minimums = [(0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573)] _fmin: float = -3.32237 _fmax = 0.0 - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - _alpha: np.ndarray = np.array([1.0, 1.2, 3.0, 3.2]) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - _A: np.ndarray = np.array( + _alpha: npt.NDArray = np.array([1.0, 1.2, 3.0, 3.2]) + _A: npt.NDArray = np.array( [ [10, 3, 17, 3.5, 1.7, 8], [0.05, 10, 17, 0.1, 8, 14], @@ -209,8 +202,7 @@ class Hartmann6(SyntheticFunction): [17, 8, 0.05, 10, 0.1, 14], ] ) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - _P: np.ndarray = 10 ** (-4) * np.array( + _P: npt.NDArray = 10 ** (-4) * np.array( [ [1312, 1696, 5569, 124, 8283, 5886], [2329, 4135, 8307, 3736, 1004, 9991], @@ -221,8 +213,7 @@ class Hartmann6(SyntheticFunction): @override @copy_doc(SyntheticFunction._f) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _f(self, X: np.ndarray) -> float: + def _f(self, X: npt.NDArray) -> float: y = 0.0 for j, alpha_j in enumerate(self._alpha): t = 0 @@ -245,8 +236,7 @@ class Aug_Hartmann6(Hartmann6): @override @copy_doc(SyntheticFunction._f) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _f(self, X: np.ndarray) -> float: + def _f(self, X: npt.NDArray) -> float: y = 0.0 alpha_0 = self._alpha[0] - 0.1 * (1 - X[-1]) for j, alpha_j in enumerate(self._alpha): @@ -275,8 +265,7 @@ class Branin(SyntheticFunction): @override @copy_doc(SyntheticFunction._f) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _f(self, X: np.ndarray) -> float: + def _f(self, X: npt.NDArray) -> float: x_1 = X[0] x_2 = X[1] return float( @@ -301,8 +290,7 @@ class Aug_Branin(SyntheticFunction): @override @copy_doc(SyntheticFunction._f) - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - def _f(self, X: np.ndarray) -> float: + def _f(self, X: npt.NDArray) -> float: x_1 = X[0] x_2 = X[1] return float( diff --git a/ax/utils/measurement/tests/test_synthetic_functions.py b/ax/utils/measurement/tests/test_synthetic_functions.py index be9c1708755..efaaa3ce9f3 100644 --- a/ax/utils/measurement/tests/test_synthetic_functions.py +++ b/ax/utils/measurement/tests/test_synthetic_functions.py @@ -77,8 +77,6 @@ def test_aug_hartmann6(self) -> None: # SupportsAbs[SupportsRound[object]]]` but got `float`. self.assertAlmostEqual(aug_hartmann6(1, 2, 3, 4, 5, 6, 1), 0.0) self.assertAlmostEqual( - # pyre-fixme[6]: For 2nd argument expected `SupportsRSub[Variable[_T], - # SupportsAbs[SupportsRound[object]]]` but got `float`. aug_hartmann6(x1=1, x2=2, x3=3, x4=4, x5=5, x6=6, x7=1), # pyre-fixme[6]: For 2nd argument expected `SupportsRSub[Variable[_T], # SupportsAbs[SupportsRound[object]]]` but got `float`. diff --git a/ax/utils/sensitivity/sobol_measures.py b/ax/utils/sensitivity/sobol_measures.py index c64d9f2da1c..d9e598b31c9 100644 --- a/ax/utils/sensitivity/sobol_measures.py +++ b/ax/utils/sensitivity/sobol_measures.py @@ -10,7 +10,7 @@ from copy import deepcopy from typing import Any -import numpy as np +import numpy.typing as npt import torch from ax.modelbridge.torch import TorchModelBridge from ax.models.torch.botorch import BotorchModel @@ -830,8 +830,7 @@ def ax_parameter_sens( order: str = "first", signed: bool = True, **sobol_kwargs: Any, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> dict[str, dict[str, np.ndarray]]: +) -> dict[str, dict[str, npt.NDArray]]: """ Compute sensitivity for all metrics on an TorchModelBridge. @@ -984,13 +983,10 @@ def _get_model_per_metric( def array_with_string_indices_to_dict( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. rows: list[str], cols: list[str], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - A: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> dict[str, dict[str, np.ndarray]]: + A: npt.NDArray, +) -> dict[str, dict[str, npt.NDArray]]: """ Args: rows: A list of strings with which to index rows of A. diff --git a/ax/utils/sensitivity/tests/test_sensitivity.py b/ax/utils/sensitivity/tests/test_sensitivity.py index 320a1fe659a..e1f45da844e 100644 --- a/ax/utils/sensitivity/tests/test_sensitivity.py +++ b/ax/utils/sensitivity/tests/test_sensitivity.py @@ -301,10 +301,6 @@ def test_SobolGPMean_SAASBO(self) -> None: for i, row in enumerate(ind_dict): for j, col in enumerate(ind_dict[row]): self.assertAlmostEqual( - # pyre-fixme[6]: For 2nd argument expected - # `SupportsRSub[Variable[_T], - # SupportsAbs[SupportsRound[object]]]` but got - # `Union[bool, float, int]`. ind_dict[row][col], # pyre-fixme[6]: For 2nd argument expected # `SupportsRSub[Variable[_T], @@ -337,9 +333,6 @@ def test_SobolGPMean_SAASBO(self) -> None: ) # check that the first and second order indices are the same self.assertAlmostEqual( - # pyre-fixme[6]: For 2nd argument expected - # `SupportsRSub[Variable[_T], SupportsAbs[SupportsRound[object]]]` - # but got `Union[bool, float, int]`. second_ind_dict["branin"]["x1"], # pyre-fixme[6]: For 2nd argument expected # `SupportsRSub[Variable[_T], SupportsAbs[SupportsRound[object]]]` @@ -347,9 +340,6 @@ def test_SobolGPMean_SAASBO(self) -> None: fo_ind_tnsr[0, 0].item(), ) self.assertAlmostEqual( - # pyre-fixme[6]: For 2nd argument expected - # `SupportsRSub[Variable[_T], SupportsAbs[SupportsRound[object]]]` - # but got `Union[bool, float, int]`. second_ind_dict["branin"]["x2"], # pyre-fixme[6]: For 2nd argument expected # `SupportsRSub[Variable[_T], SupportsAbs[SupportsRound[object]]]` @@ -357,9 +347,6 @@ def test_SobolGPMean_SAASBO(self) -> None: fo_ind_tnsr[0, -1].item(), ) self.assertAlmostEqual( - # pyre-fixme[6]: For 2nd argument expected - # `SupportsRSub[Variable[_T], SupportsAbs[SupportsRound[object]]]` - # but got `Union[bool, float, int]`. second_ind_dict["branin"]["x1 & x2"], # pyre-fixme[6]: For 2nd argument expected # `SupportsRSub[Variable[_T], SupportsAbs[SupportsRound[object]]]` diff --git a/ax/utils/stats/model_fit_stats.py b/ax/utils/stats/model_fit_stats.py index 5ff548c18a6..f3aeeaf2e4e 100644 --- a/ax/utils/stats/model_fit_stats.py +++ b/ax/utils/stats/model_fit_stats.py @@ -10,6 +10,7 @@ from typing import Protocol import numpy as np +import numpy.typing as npt from ax.utils.common.logger import get_logger from scipy.stats import fisher_exact, norm, pearsonr, spearmanr @@ -30,24 +31,17 @@ class ModelFitMetricProtocol(Protocol): """Structural type for model fit metrics.""" def __call__( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. self, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: ... def compute_model_fit_metrics( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: Mapping[str, np.ndarray], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: Mapping[str, np.ndarray], - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: Mapping[str, np.ndarray], + y_obs: Mapping[str, npt.NDArray], + y_pred: Mapping[str, npt.NDArray], + se_pred: Mapping[str, npt.NDArray], fit_metrics_dict: Mapping[str, ModelFitMetricProtocol], ) -> dict[str, dict[str, float]]: """Computes the model fit metrics for each experimental metric in the input dicts. @@ -79,12 +73,9 @@ def compute_model_fit_metrics( def coefficient_of_determination( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray | None = None, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray | None = None, eps: float = 1e-12, ) -> float: """Computes coefficient of determination, the proportion of variance in `y_obs` @@ -105,12 +96,9 @@ def coefficient_of_determination( def mean_of_the_standardized_error( # i.e. standardized bias - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: """Computes the mean of the error standardized by the predictive standard deviation of the model `se_pred`. If the model makes good predictions and its uncertainty is @@ -132,12 +120,9 @@ def mean_of_the_standardized_error( # i.e. standardized bias def std_of_the_standardized_error( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: """Standard deviation of the error standardized by the predictive standard deviation of the model `se_pred`. If the uncertainty is quantified well, should be close to 1. @@ -158,12 +143,9 @@ def std_of_the_standardized_error( def entropy_of_observations( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, bandwidth: float = DEFAULT_KDE_BANDWIDTH, ) -> float: """Computes the entropy of the observations y_obs using a kernel density estimator. @@ -198,8 +180,7 @@ def entropy_of_observations( return _entropy_via_kde(y_obs, bandwidth=bandwidth) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _entropy_via_kde(y: np.ndarray, bandwidth: float = DEFAULT_KDE_BANDWIDTH) -> float: +def _entropy_via_kde(y: npt.NDArray, bandwidth: float = DEFAULT_KDE_BANDWIDTH) -> float: """Computes the entropy of the kernel density estimate of the input data. Args: @@ -216,67 +197,52 @@ def _entropy_via_kde(y: np.ndarray, bandwidth: float = DEFAULT_KDE_BANDWIDTH) -> def _mean_prediction_ci( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: # Pyre does not allow float * np.ndarray. return float(np.mean(1.96 * 2 * se_pred / np.abs(y_obs))) def _log_likelihood( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: return float(np.sum(norm.logpdf(y_obs, loc=y_pred, scale=se_pred))) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _mape(y_obs: np.ndarray, y_pred: np.ndarray, se_pred: np.ndarray) -> float: +def _mape(y_obs: npt.NDArray, y_pred: npt.NDArray, se_pred: npt.NDArray) -> float: """Mean absolute predictive error""" eps = np.finfo(y_obs.dtype).eps return float(np.mean(np.abs(y_pred - y_obs) / np.abs(y_obs).clip(min=eps))) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _mse(y_obs: np.ndarray, y_pred: np.ndarray, se_pred: np.ndarray) -> float: +def _mse(y_obs: npt.NDArray, y_pred: npt.NDArray, se_pred: npt.NDArray) -> float: """Mean squared error""" return float(np.mean((y_pred - y_obs) ** 2)) -# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -def _wmape(y_obs: np.ndarray, y_pred: np.ndarray, se_pred: np.ndarray) -> float: +def _wmape(y_obs: npt.NDArray, y_pred: npt.NDArray, se_pred: npt.NDArray) -> float: """Weighted mean absolute predictive error""" eps = np.finfo(y_obs.dtype).eps return float(np.sum(np.abs(y_pred - y_obs)) / np.sum(np.abs(y_obs)).clip(min=eps)) def _total_raw_effect( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: min_y_obs = np.min(y_obs) return float((np.max(y_obs) - min_y_obs) / min_y_obs) def _correlation_coefficient( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: with np.errstate(invalid="ignore"): rho, _ = pearsonr(y_pred, y_obs) @@ -284,12 +250,9 @@ def _correlation_coefficient( def _rank_correlation( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: with np.errstate(invalid="ignore"): rho, _ = spearmanr(y_pred, y_obs) @@ -297,12 +260,9 @@ def _rank_correlation( def _fisher_exact_test_p( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_obs: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - y_pred: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - se_pred: np.ndarray, + y_obs: npt.NDArray, + y_pred: npt.NDArray, + se_pred: npt.NDArray, ) -> float: """Perform a Fisher exact test on the contingency table constructed from agreement/disagreement between the predicted and observed data. diff --git a/ax/utils/stats/statstools.py b/ax/utils/stats/statstools.py index f539142e096..004b6b70b75 100644 --- a/ax/utils/stats/statstools.py +++ b/ax/utils/stats/statstools.py @@ -12,6 +12,7 @@ from typing import Union import numpy as np +import numpy.typing as npt import pandas as pd from ax.core.data import Data from ax.utils.common.logger import get_logger @@ -22,10 +23,8 @@ def inverse_variance_weight( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - means: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - variances: np.ndarray, + means: npt.NDArray, + variances: npt.NDArray, conflicting_noiseless: str = "warn", ) -> tuple[float, float]: """Perform inverse variance weighting. @@ -64,12 +63,9 @@ def inverse_variance_weight( def total_variance( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - means: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - variances: np.ndarray, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sample_sizes: np.ndarray, + means: npt.NDArray, + variances: npt.NDArray, + sample_sizes: npt.NDArray, ) -> float: """Compute total variance.""" variances = variances * sample_sizes @@ -83,8 +79,7 @@ def total_variance( def positive_part_james_stein( means: num_mixed, sems: num_mixed, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[npt.NDArray, npt.NDArray]: """Estimation method for Positive-part James-Stein estimator. This method takes a vector of K means (`y_i`) and standard errors @@ -163,19 +158,15 @@ def positive_part_james_stein( def relativize( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - means_t: np.ndarray | list[float] | float, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sems_t: np.ndarray | list[float] | float, + means_t: npt.NDArray | list[float] | float, + sems_t: npt.NDArray | list[float] | float, mean_c: float, sem_c: float, bias_correction: bool = True, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - cov_means: np.ndarray | list[float] | float = 0.0, + cov_means: npt.NDArray | list[float] | float = 0.0, as_percent: bool = False, control_as_constant: bool = False, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[npt.NDArray, npt.NDArray]: """Ratio estimator based on the delta method. This uses the delta method (i.e. a Taylor series approximation) to estimate @@ -268,19 +259,15 @@ def relativize( def unrelativize( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - means_t: np.ndarray | list[float] | float, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - sems_t: np.ndarray | list[float] | float, + means_t: npt.NDArray | list[float] | float, + sems_t: npt.NDArray | list[float] | float, mean_c: float, sem_c: float, bias_correction: bool = True, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - cov_means: np.ndarray | list[float] | float = 0.0, + cov_means: npt.NDArray | list[float] | float = 0.0, as_percent: bool = False, control_as_constant: bool = False, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[npt.NDArray, npt.NDArray]: """ Reverse operation of ax.utils.stats.statstools.relativize. @@ -342,14 +329,11 @@ def unrelativize( def agresti_coull_sem( - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - n_numer: pd.Series | np.ndarray | int, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. - n_denom: pd.Series | np.ndarray | int, + n_numer: pd.Series | npt.NDArray | int, + n_denom: pd.Series | npt.NDArray | int, prior_successes: int = 2, prior_failures: int = 2, - # pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters. -) -> np.ndarray | float: +) -> npt.NDArray | float: """Compute the Agresti-Coull style standard error for a binomial proportion. Reference: diff --git a/ax/utils/testing/mock.py b/ax/utils/testing/mock.py index d4183162838..1b509af7da2 100644 --- a/ax/utils/testing/mock.py +++ b/ax/utils/testing/mock.py @@ -33,9 +33,7 @@ def fast_botorch_optimize_context_manager( USE RESPONSIBLY. """ - def one_iteration_minimize( - *args: Any, **kwargs: Any - ) -> OptimizeResult: # pyre-ignore[11] + def one_iteration_minimize(*args: Any, **kwargs: Any) -> OptimizeResult: if kwargs["options"] is None: kwargs["options"] = {}