Replace # pytype: disable supprressions with # pyrefly: ignore

PiperOrigin-RevId: 993590179

Author: oprypin

demos/run_vizier_server.py

@@ -59,7 +59,7 @@ def main(argv: Sequence[str]) -> None:
     raise app.UsageError('Too many command-line arguments.')
 
   server = servers.DefaultVizierServer(
-      host=FLAGS.host, database_url=FLAGS.database_url  # pyrefly: ignore[unexpected-keyword]
+      host=FLAGS.host, database_url=FLAGS.database_url
   )
   logging.info('Address to Vizier Server is: %s', server.endpoint)
 

vizier/_src/algorithms/classification/classifiers.py

@@ -92,4 +92,4 @@ def __call__(self) -> np.ndarray:
     if self.eval_metric == 'probability':
       return self.classifier.predict_proba(np.asarray(self.features_test))[:, 1]
     else:
-      return self.classifier.decision_function(np.asarray(self.features_test))  # pytype:disable=attribute-error
+      return self.classifier.decision_function(np.asarray(self.features_test))  # pyrefly: ignore[missing-attribute]

vizier/_src/algorithms/designers/eagle_strategy/eagle_strategy.py

@@ -146,8 +146,8 @@ def __init__(
     self._rng = np.random.default_rng(seed=seed)
     self._config = config or FireflyAlgorithmConfig()
     self._utils = EagleStrategyUtils(self._problem, self._config, self._rng)
-    self._firefly_pool = FireflyPool(  # pyrefly: ignore[missing-argument]
-        utils=self._utils, capacity=self._utils.compute_pool_capacity()  # pyrefly: ignore[unexpected-keyword]
+    self._firefly_pool = FireflyPool(
+        utils=self._utils, capacity=self._utils.compute_pool_capacity()
     )
 
     if initial_designer_factory is None:

vizier/_src/algorithms/designers/eagle_strategy/eagle_strategy_utils.py

@@ -361,7 +361,7 @@ def get_metric(self, trial: vz.Trial) -> float:
       return np.nan
     if trial.final_measurement is None:
       raise ValueError('Trial is not completed.')
-    return trial.final_measurement.metrics[OBJECTIVE_NAME]  # pytype: disable=bad-return-type
+    return trial.final_measurement.metrics[OBJECTIVE_NAME]
 
   def is_better_than(
       self,

vizier/_src/algorithms/designers/gp/gp_models_test.py

@@ -66,7 +66,7 @@ def _setup_lambda_search(
   problem = vz.ProblemStatement(
       search_space=search_space,
       metric_information=vz.MetricsConfig(
-          metrics=[  # pyrefly: ignore[unexpected-keyword]
+          metrics=[
               vz.MetricInformation('obj', goal=vz.ObjectiveMetricGoal.MAXIMIZE),
           ]
       ),
@@ -378,7 +378,7 @@ def test_multi_task(self, multitask_type: mt_type):
     problem = vz.ProblemStatement(
         search_space=search_space,
         metric_information=vz.MetricsConfig(
-            metrics=[  # pyrefly: ignore[unexpected-keyword]
+            metrics=[
                 vz.MetricInformation(
                     'obj1', goal=vz.ObjectiveMetricGoal.MAXIMIZE
                 ),

vizier/_src/algorithms/designers/gp/output_warpers.py

@@ -357,10 +357,10 @@ def warp(self, labels_arr: types.Array) -> types.Array:
         labels_arr[i] = rank_ppf * estimated_std + median
 
     # Save information needed for unwarping.
-    self._unwarper = _HalfRankUnwarper(  # pyrefly: ignore[missing-argument]
-        original_labels=unique_labels,  # pyrefly: ignore[unexpected-keyword]
-        warped_labels=labels_arr[is_finite][unique_idx],  # pyrefly: ignore[bad-argument-type, unexpected-keyword]
-        original_label_median=unique_labels[len(unique_labels) // 2],  # pyrefly: ignore[unexpected-keyword]
+    self._unwarper = _HalfRankUnwarper(
+        original_labels=unique_labels,
+        warped_labels=labels_arr[is_finite][unique_idx],  # pyrefly: ignore[bad-argument-type]
+        original_label_median=unique_labels[len(unique_labels) // 2],
     )
     return labels_arr[:, np.newaxis]
 

vizier/_src/algorithms/designers/meta_learning/meta_learning.py

@@ -157,11 +157,11 @@ def __attrs_post_init__(self):
       self.seed = np.random.randint(low=0, high=1e6)  # pyrefly: ignore[no-matching-overload]
 
     # Instantiate an MetaLearningUtils.
-    self._utils = utils.MetaLearningUtils(  # pyrefly: ignore[missing-argument]
-        goal=self.problem.metric_information.item().goal,  # pyrefly: ignore[unexpected-keyword]
-        tuned_metric_name=self.problem.metric_information.item().name,  # pyrefly: ignore[unexpected-keyword]
-        meta_metric_name=self._meta_designer_metric_name,  # pyrefly: ignore[unexpected-keyword]
-        tuning_params=self.tuning_hyperparams,  # pyrefly: ignore[unexpected-keyword]
+    self._utils = utils.MetaLearningUtils(
+        goal=self.problem.metric_information.item().goal,
+        tuned_metric_name=self.problem.metric_information.item().name,
+        meta_metric_name=self._meta_designer_metric_name,
+        tuning_params=self.tuning_hyperparams,
     )
     # Instantiated 'tuned' designer the with default hyper-parameters.
     self._curr_tuned_hyperparams = self._utils.get_default_hyperparameters()

vizier/_src/algorithms/designers/meta_learning/meta_learning_utils.py

@@ -38,7 +38,7 @@ def meta_problem(self) -> vz.ProblemStatement:
     """Create meta problem."""
     problem = vz.ProblemStatement(search_space=self._tuning_params)
     problem.metric_information = vz.MetricsConfig(
-        metrics=[vz.MetricInformation(self._meta_metric_name, goal=self._goal)]  # pyrefly: ignore[unexpected-keyword]
+        metrics=[vz.MetricInformation(self._meta_metric_name, goal=self._goal)]
     )
     return problem
 

vizier/_src/algorithms/designers/meta_learning/meta_learning_utils_test.py

@@ -29,11 +29,11 @@ def setUp(self):
     super().setUp()
     space = vz.SearchSpace()
     space.root.add_int_param('tuned_param', 0, 100, default_value=55)
-    self.utils = meta_learning_utils.MetaLearningUtils(  # pyrefly: ignore[missing-argument]
-        goal=vz.ObjectiveMetricGoal.MAXIMIZE,  # pyrefly: ignore[unexpected-keyword]
-        tuned_metric_name='tuned_obj',  # pyrefly: ignore[unexpected-keyword]
-        meta_metric_name='meta_obj',  # pyrefly: ignore[unexpected-keyword]
-        tuning_params=space,  # pyrefly: ignore[unexpected-keyword]
+    self.utils = meta_learning_utils.MetaLearningUtils(
+        goal=vz.ObjectiveMetricGoal.MAXIMIZE,
+        tuned_metric_name='tuned_obj',
+        meta_metric_name='meta_obj',
+        tuning_params=space,
     )
     self.meta_trials = []
     for i in range(10):

vizier/_src/algorithms/ensemble/ensemble_design.py

@@ -97,7 +97,6 @@ def update(self, observation: IndexWithReward):  # pyrefly: ignore[bad-override]
     self._history.append(observation)
 
 
-# pytype: disable=attribute-error
 # https://www.cs.princeton.edu/courses/archive/fall16/cos402/lectures/402-lec22.pdf.
 @attrs.define
 class EXP3UniformEnsembleDesign(EnsembleDesign):
@@ -129,7 +128,7 @@ def ensemble_probs(self) -> np.ndarray:
     probs = (1 - gamma) * softmax(self._log_weights) + gamma * uniform
     return probs
 
-  def update(self, observation: IndexWithReward):
+  def update(self, observation: IndexWithReward):  # pyrefly: ignore[bad-override]
     """Update history and weights."""
     expert_idx, reward = observation
     if not self.use_reward_estimator:
@@ -242,7 +241,7 @@ def observation_probs(self) -> np.ndarray:
     ) + algo_prob_sum / (2 * np.sum(algo_prob_sum))
     return observation_prob
 
-  def update(self, observation: IndexWithReward):
+  def update(self, observation: IndexWithReward):  # pyrefly: ignore[bad-override]
     expert_idx, reward = observation
     reward = min(reward, self.max_reward)
     reward_estimator = reward * 1.0 / self.ensemble_probs[expert_idx]
@@ -278,6 +277,3 @@ def update(self, observation: IndexWithReward):
       base_algo.update((expert_idx, reward_estimator))
 
     self._history.append(observation)
-
-
-# pytype: enable=attribute-error