Replace # pytype: disable suppressions with # pyrefly: ignore

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

vizier/_src/algorithms/evolution/nsga2.py

@@ -86,19 +86,19 @@ def crowding_distance(
 
     # Compute the range of the m-th metric.
     yy = ys[:, m]  # Shape: (num_population,)
-    yrange = yy[sid[-1]] - yy[sid[0]] + np.finfo(np.float32).eps  # pyrefly: ignore[bad-index]
+    yrange = yy[sid[-1]] - yy[sid[0]] + np.finfo(np.float32).eps
 
     # Lower boundary is assigned a one-sided score and does not automatically
     # get infinity. This is different from the paper. The lower boundary means
     # it's dominated by all other points in one dimension. There's no reason to
     # favor it over other points.
-    scores[sid[0]] += (yy[sid[1]] - yy[sid[0]]) / yrange  # pyrefly: ignore[bad-index, unsupported-operation]
+    scores[sid[0]] += (yy[sid[1]] - yy[sid[0]]) / yrange
     # Upper boundary is assigned infinity. This point will survive anyways
     # because it's pareto-optimal. But in case there are ties, it's useful to
     # make only one of them stand out.
-    scores[sid[-1]] += np.inf  # pyrefly: ignore[bad-index, unsupported-operation]
+    scores[sid[-1]] += np.inf
 
-    scores[sid[1:-1]] += (yy[sid[2:]] - yy[sid[:-2]]) / yrange  # pyrefly: ignore[bad-index, unsupported-operation]
+    scores[sid[1:-1]] += (yy[sid[2:]] - yy[sid[:-2]]) / yrange
   # Normalize the score to [0, 1].
   return scores / ys.shape[1]
 

vizier/_src/algorithms/evolution/numpy_populations.py

@@ -134,7 +134,7 @@ def __init__(
     if ids is None:
       ids = np.zeros([xs.shape[0]])
 
-    self.__attrs_init__(xs, ids, generations)  # pyrefly: ignore[missing-attribute]
+    self.__attrs_init__(xs, ids, generations)
 
   def __len__(self) -> int:
     return self.generations.shape[0]
@@ -325,7 +325,7 @@ def __init__(
         self._trial_converter.to_features([])
     )
 
-  def to_suggestions(  # pytype: disable=signature-mismatch  # overriding-parameter-type-checks
+  def to_suggestions(  # pyrefly: ignore[bad-override]
       self, offsprings: Offspring
   ) -> Collection[vz.TrialSuggestion]:
     parameters_list = self._trial_converter.to_parameters(

vizier/_src/algorithms/policies/designer_policy.py

@@ -181,7 +181,7 @@ def __init__(
     self._ns_root = ns_root
     self._cache: trial_caches.IdDeduplicatingTrialLoader = (
         trial_caches.IdDeduplicatingTrialLoader(
-            supporter, include_intermediate_measurements=False  # pyrefly: ignore[unexpected-keyword]
+            supporter, include_intermediate_measurements=False
         )
     )
     self._problem_statement = problem_statement

vizier/_src/algorithms/policies/random_policy_test.py

@@ -42,8 +42,8 @@ def test_make_suggestions(self):
     num_suggestions = 5
     num_params = len(self.study_config.search_space.parameters)
 
-    suggest_request = pythia.SuggestRequest(  # pyrefly: ignore[missing-argument]
-        study_descriptor=self.policy_supporter.study_descriptor(),  # pyrefly: ignore[unexpected-keyword]
+    suggest_request = pythia.SuggestRequest(
+        study_descriptor=self.policy_supporter.study_descriptor(),
         count=num_suggestions,
     )
     decisions = self.policy.suggest(suggest_request)
@@ -60,8 +60,8 @@ def test_make_early_stopping_decisions(self):
     request_trial_ids = [1, 2]
     trial_ids_stopped = set()
     for _ in range(count):
-      request = pythia.EarlyStopRequest(  # pyrefly: ignore[missing-argument]
-          study_descriptor=self.policy_supporter.study_descriptor(),  # pyrefly: ignore[unexpected-keyword]
+      request = pythia.EarlyStopRequest(
+          study_descriptor=self.policy_supporter.study_descriptor(),
           trial_ids=request_trial_ids,
       )
       early_stop_decisions = self.policy.early_stop(request)

vizier/_src/algorithms/testing/comparator_runner.py

@@ -104,8 +104,8 @@ def assert_better_efficiency(
     candidate_curve = analyzers.ConvergenceCurve.align_xs(
         candidate_curves, interpolate_repeats=True
     )[0]
-    comparator = analyzers.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=baseline_curve, compared_curve=candidate_curve  # pyrefly: ignore[unexpected-keyword]
+    comparator = analyzers.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=baseline_curve, compared_curve=candidate_curve
     )
 
     if (log_eff_score := comparator.score()) < score_threshold:
@@ -167,14 +167,14 @@ def assert_optimizer_better_simple_regret(
     candidate_optimizer = candidate_optimizer_factory(converter)
 
     for i in range(self.baseline_num_repeats):
-      res = baseline_optimizer(score_fn, count=1, seed=random.PRNGKey(i))  # pytype: disable=wrong-arg-types
+      res = baseline_optimizer(score_fn, count=1, seed=random.PRNGKey(i))  # pyrefly: ignore[bad-argument-type]
       trial = vb.best_candidates_to_trials(res, converter)
       baseline_obj_values.append(
           trial[0].final_measurement_or_die.metrics['acquisition'].value
       )
 
     for i in range(self.candidate_num_repeats):
-      res = candidate_optimizer(score_fn, count=1, seed=random.PRNGKey(i))  # pytype: disable=wrong-arg-types
+      res = candidate_optimizer(score_fn, count=1, seed=random.PRNGKey(i))  # pyrefly: ignore[bad-argument-type]
       trial = vb.best_candidates_to_trials(res, converter)
       candidate_obj_values.append(
           trial[0].final_measurement_or_die.metrics['acquisition'].value

vizier/_src/benchmarks/experimenters/switch_experimenter_test.py

@@ -37,8 +37,8 @@ def test_switch_experimenter(self):
 
     switch_exptr.evaluate([t0, t1])
 
-    self.assertEqual(t0.final_measurement.metrics['switch_metric'].value, 0.0)  # pytype:disable=attribute-error
-    self.assertEqual(t1.final_measurement.metrics['switch_metric'].value, 100.0)  # pytype:disable=attribute-error
+    self.assertEqual(t0.final_measurement.metrics['switch_metric'].value, 0.0)  # pyrefly: ignore[missing-attribute]
+    self.assertEqual(t1.final_measurement.metrics['switch_metric'].value, 100.0)  # pyrefly: ignore[missing-attribute]
 
 
 if __name__ == '__main__':

vizier/_src/benchmarks/runners/benchmark_state.py

@@ -85,7 +85,7 @@ def from_designer_factory(
         designer_factory=designer_factory,
         seed=seed,
     )
-    return PolicySuggester(policy=policy, local_supporter=supporter)  # pyrefly: ignore[missing-argument, unexpected-keyword]
+    return PolicySuggester(policy=policy, local_supporter=supporter)
 
 
 @attr.define
@@ -162,8 +162,8 @@ def __call__(self, seed: Optional[int] = None) -> BenchmarkState:
     problem = self.experimenter.problem_statement()
     return BenchmarkState(
         experimenter=self.experimenter,
-        algorithm=PolicySuggester(  # pyrefly: ignore[missing-argument]
-            policy=self.policy_factory(problem, seed),  # pyrefly: ignore[bad-argument-type, unexpected-keyword]
-            local_supporter=pythia.InRamPolicySupporter(problem),  # pyrefly: ignore[unexpected-keyword]
+        algorithm=PolicySuggester(
+            policy=self.policy_factory(problem, seed),  # pyrefly: ignore[bad-argument-type]
+            local_supporter=pythia.InRamPolicySupporter(problem),
         ),
     )

vizier/_src/jax/models/gaussian_process_model.py

@@ -120,33 +120,34 @@ def __call__(  # pyrefly: ignore[bad-override]
     Returns:
       A `tfd.GaussianProcess` with the given index points.
     """
-    # pytype: disable=not-callable  # jnp-type
     amplitude = yield sp_model.ModelParameter.from_prior(
-        tfd.LogNormal(self.dtype(0.0), 1.0, name='amplitude'),
-        constraint=sp_model.Constraint(bounds=(self.dtype(0.0), None)),
+        tfd.LogNormal(self.dtype(0.0), 1.0, name='amplitude'),  # pyrefly: ignore[not-callable]
+        constraint=sp_model.Constraint(bounds=(self.dtype(0.0), None)),  # pyrefly: ignore[not-callable]
     )
     kernel = self._kernel_class(
         amplitude=amplitude,
-        length_scale=self.dtype(1.0),
+        length_scale=self.dtype(1.0),  # pyrefly: ignore[not-callable]
         validate_args=self._use_tfp_runtime_validation,
     )
     inverse_length_scale_continuous = yield sp_model.ModelParameter.from_prior(
         tfd.Sample(
-            tfd.LogNormal(self.dtype(0.0), 1.0),
+            tfd.LogNormal(self.dtype(0.0), 1.0),  # pyrefly: ignore[not-callable]
             sample_shape=(self.dimension.continuous,),
             name='inverse_length_scale_continuous',
         ),
-        constraint=sp_model.Constraint(bounds=(self.dtype(0.0), None)),
+        constraint=sp_model.Constraint(bounds=(self.dtype(0.0), None)),  # pyrefly: ignore[not-callable]
     )
     inverse_length_scale_categorical = yield sp_model.ModelParameter.from_prior(
         tfd.Sample(
-            tfd.LogNormal(self.dtype(0.0), 1.0),
+            tfd.LogNormal(self.dtype(0.0), 1.0),  # pyrefly: ignore[not-callable]
             sample_shape=(self.dimension.categorical,),
             name='inverse_length_scale_categorical',
         ),
         constraint=sp_model.Constraint(
-            bounds=(jnp.zeros(self.dimension.categorical, dtype=self.dtype),
-                    None)
+            bounds=(
+                jnp.zeros(self.dimension.categorical, dtype=self.dtype),
+                None,
+            )
         ),
     )
     kernel = tfpke.FeatureScaledWithCategorical(
@@ -157,11 +158,13 @@ def __call__(  # pyrefly: ignore[bad-override]
         validate_args=self._use_tfp_runtime_validation,
     )
     observation_noise_variance = yield sp_model.ModelParameter.from_prior(
-        tfd.LogNormal(self.dtype(0.0), 1.0,
-                      name='observation_noise_variance'),
-        constraint=sp_model.Constraint(bounds=(self.dtype(0.0), None)),
+        tfd.LogNormal(
+            self.dtype(0.0),  # pyrefly: ignore[not-callable]
+            1.0,
+            name='observation_noise_variance',
+        ),
+        constraint=sp_model.Constraint(bounds=(self.dtype(0.0), None)),  # pyrefly: ignore[not-callable]
     )
-    # pytype: enable=not-callable
     if inputs is not None:
       inputs = tfpke.ContinuousAndCategoricalValues(
           inputs.continuous.padded_array, inputs.categorical.padded_array

vizier/_src/jax/optimizers/jaxopt_wrappers_test.py

@@ -53,8 +53,8 @@ def test_sinusoidal(self, cls, bounds, nest_constraint):
 
   def test_max_duration(self):
     optimizer = jw.JaxoptScipyLbfgsB(
-        max_duration=datetime.timedelta(seconds=0),  # pyrefly: ignore[unexpected-keyword]
-        speed_test=True,  # pyrefly: ignore[unexpected-keyword]
+        max_duration=datetime.timedelta(seconds=0),
+        speed_test=True,
     )
     random_restarts = 3
     rngs = jax.random.split(jax.random.PRNGKey(1), random_restarts + 1)

vizier/_src/jax/xla_pareto.py

@@ -90,11 +90,11 @@ def is_frontier(
   frontier = np.ones(ys.shape[0], dtype=np.bool_)
 
   for begin, end in zip(idx[1:], idx[:-1]):
-    candidates = ys[frontier]  # pyrefly: ignore[bad-index]
+    candidates = ys[frontier]
     # Filter candidates by comparing against the slice.
     if verbose:
       print(f"Compare {len(candidates)} against {begin}:{end}.")
-    tt = _is_pareto_optimal_against(candidates, ys[begin:end], strict=True)  # pyrefly: ignore[bad-index]
+    tt = _is_pareto_optimal_against(candidates, ys[begin:end], strict=True)
     frontier[frontier] = tt
   return frontier
 
@@ -148,7 +148,7 @@ def get_frontier(
     if verbose:
       # Use print. This method won't run in production anyways.
       print(f"Compare {len(candidates)} against {begin}:{end}.")
-    tt = _is_pareto_optimal_against(candidates, ys[begin:end], strict=True)  # pyrefly: ignore[bad-index]
+    tt = _is_pareto_optimal_against(candidates, ys[begin:end], strict=True)
     candidates = candidates[tt]
   return candidates
 

vizier/_src/pythia/local_policy_supporters.py

@@ -276,8 +276,10 @@ def SuggestTrials(self, algorithm: policy.Policy,
     """Suggest and add new trials."""
 
     decisions = algorithm.suggest(
-        policy.SuggestRequest(  # pyrefly: ignore[missing-argument]
-            study_descriptor=self.study_descriptor(), count=count))  # pyrefly: ignore[unexpected-keyword]
+        policy.SuggestRequest(
+            study_descriptor=self.study_descriptor(), count=count
+        )
+    )
     self._UpdateMetadata(decisions.metadata)
     return self.AddSuggestions([
         vz.TrialSuggestion(d.parameters, metadata=d.metadata)