diff --git a/vizier/_src/algorithms/optimizers/eagle_strategy.py b/vizier/_src/algorithms/optimizers/eagle_strategy.py
index a3f72f67e..3d1328e26 100644
--- a/vizier/_src/algorithms/optimizers/eagle_strategy.py
+++ b/vizier/_src/algorithms/optimizers/eagle_strategy.py
@@ -389,7 +389,6 @@ def __call__(
       # This configuration updates all the fireflies in each iteration.
       suggestion_batch_size = pool_size
     # Use priors to populate Eagle state
-    # pytype: disable=wrong-arg-types  # jnp-type
     return VectorizedEagleStrategy(
         n_feature_dimensions=feature_dimensions.n_feature_dimensions,
         n_feature_dimensions_with_padding=(
@@ -402,9 +401,8 @@ def __call__(
         pool_size=pool_size,
         categorical_sizes=jnp.array(feature_dimensions.categorical_sizes),
         max_categorical_size=feature_dimensions.max_categorical_size,
-        dtype=converter._impl.dtype,
+        dtype=converter._impl.dtype,  # pyrefly: ignore[bad-argument-type]
     )
-    # pytype: enable=wrong-arg-types
 
 
 @struct.dataclass
@@ -555,15 +553,13 @@ def init_state(
           n_parallel=n_parallel,
           seed=seed,
       )
-    # pytype: disable=wrong-arg-types  # jnp-type
     return VectorizedEagleStrategyState(
         iterations=jnp.array(0),
         features=init_features,
         rewards=jnp.ones(self.pool_size) * -jnp.inf,
-        best_reward=-jnp.inf,
+        best_reward=-jnp.inf,  # pyrefly: ignore[bad-argument-type]
         perturbations=jnp.ones(self.pool_size) * self.config.perturbation,
     )
-    # pytype: enable=wrong-arg-types
 
   def _populate_pool_with_prior_trials(
       self,
@@ -847,14 +843,12 @@ def _create_features(
     if self.config.mutate_normalization_type == MutateNormalizationType.MEAN:
       # Divide the push / pull forces by the number of participating fireflies.
       # Also multiply by normalization_scale.
-      # pytype: disable=wrong-arg-types  # jnp-type
       norm_scaled_pulls = self.config.normalization_scale * jnp.nan_to_num(
-          scaled_pulls / jnp.sum(scaled_pulls > 0.0, axis=1, keepdims=True), 0
+          scaled_pulls / jnp.sum(scaled_pulls > 0.0, axis=1, keepdims=True), 0  # pyrefly: ignore[bad-argument-type]
       )
       norm_scaled_push = self.config.normalization_scale * jnp.nan_to_num(
-          scaled_push / jnp.sum(scaled_push < 0.0, axis=1, keepdims=True), 0
+          scaled_push / jnp.sum(scaled_push < 0.0, axis=1, keepdims=True), 0  # pyrefly: ignore[bad-argument-type]
       )
-      # pytype: enable=wrong-arg-types
     elif self.config.mutate_normalization_type == (
         MutateNormalizationType.RANDOM
     ):
diff --git a/vizier/_src/algorithms/optimizers/eagle_strategy_test.py b/vizier/_src/algorithms/optimizers/eagle_strategy_test.py
index 04512f5b5..5010ede49 100644
--- a/vizier/_src/algorithms/optimizers/eagle_strategy_test.py
+++ b/vizier/_src/algorithms/optimizers/eagle_strategy_test.py
@@ -548,8 +548,8 @@ def test_optimize_with_eagle_padding(self):
     converter = converters.TrialToModelInputConverter.from_problem(
         problem,
         padding_schedule=padding.PaddingSchedule(
-            num_trials=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
-            num_features=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
+            num_trials=padding.PaddingType.POWERS_OF_2,
+            num_features=padding.PaddingType.POWERS_OF_2,
         ),
     )
     eagle_factory = eagle_strategy.VectorizedEagleStrategyFactory()
@@ -575,8 +575,8 @@ def test_singleton_constraints_are_respected_with_padding(self):
     converter = converters.TrialToModelInputConverter.from_problem(
         problem,
         padding_schedule=padding.PaddingSchedule(
-            num_trials=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
-            num_features=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
+            num_trials=padding.PaddingType.POWERS_OF_2,
+            num_features=padding.PaddingType.POWERS_OF_2,
         ),
     )
     eagle_factory = eagle_strategy.VectorizedEagleStrategyFactory()
@@ -618,8 +618,8 @@ def test_compute_feature_dimensions_from_converter(
     converter = converters.TrialToModelInputConverter.from_problem(
         problem,
         padding_schedule=padding.PaddingSchedule(
-            num_trials=padding.PaddingType.MULTIPLES_OF_10,  # pyrefly: ignore[unexpected-keyword]
-            num_features=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
+            num_trials=padding.PaddingType.MULTIPLES_OF_10,
+            num_features=padding.PaddingType.POWERS_OF_2,
         ),
     )
     feature_dimensions = (
@@ -685,7 +685,7 @@ def score_fn(x, _):
     )
 
     padding_schedule = padding.PaddingSchedule(
-        num_trials=padding.PaddingType.MULTIPLES_OF_10,  # pyrefly: ignore[unexpected-keyword]
+        num_trials=padding.PaddingType.MULTIPLES_OF_10,
     )
     padding_converter = converters.TrialToModelInputConverter.from_problem(
         problem,
diff --git a/vizier/_src/algorithms/optimizers/vectorized_base.py b/vizier/_src/algorithms/optimizers/vectorized_base.py
index 0375b3601..57f209beb 100644
--- a/vizier/_src/algorithms/optimizers/vectorized_base.py
+++ b/vizier/_src/algorithms/optimizers/vectorized_base.py
@@ -313,7 +313,7 @@ class VectorizedOptimizer(Generic[_S]):
   max_evaluations: int = struct.field(pytree_node=False, default=75_000)
   dtype: types.ContinuousAndCategorical[jnp.dtype] = struct.field(
       pytree_node=False,
-      default=types.ContinuousAndCategorical[jnp.dtype](  # pytype: disable=wrong-arg-types  # jnp-type
+      default=types.ContinuousAndCategorical[jnp.dtype](
           jnp.float64, types.INT_DTYPE  # pyrefly: ignore[bad-argument-type]
       ),
   )
diff --git a/vizier/_src/benchmarks/analyzers/convergence_curve.py b/vizier/_src/benchmarks/analyzers/convergence_curve.py
index eb50a0ee5..8c886c702 100644
--- a/vizier/_src/benchmarks/analyzers/convergence_curve.py
+++ b/vizier/_src/benchmarks/analyzers/convergence_curve.py
@@ -811,16 +811,16 @@ def score(self) -> float:
     extended_compared = ConvergenceCurve.extrapolate_ys(
         compared_curve, extend_steps
     )
-    baseline_comparator = LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=combined_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=extended_baseline,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=self._baseline_quantile,  # pyrefly: ignore[unexpected-keyword]
+    baseline_comparator = LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=combined_curve,
+        compared_curve=extended_baseline,
+        compared_quantile=self._baseline_quantile,
     )
     efficiency_baseline = baseline_comparator.curve()
-    compared_comparator = LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=combined_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=extended_compared,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=self._compared_quantile,  # pyrefly: ignore[unexpected-keyword]
+    compared_comparator = LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=combined_curve,
+        compared_curve=extended_compared,
+        compared_quantile=self._compared_quantile,
     )
     efficiency_compared = compared_comparator.curve()
 
@@ -878,7 +878,7 @@ def _compute_directional_score(
         (len(compared) - i) / len(compared) if i != float('inf') else 0
         for i in convergence_curve
     ]
-    return np.mean(pct_baseline_compared)  # pyrefly: ignore[bad-return]
+    return np.mean(pct_baseline_compared)
 
   def score(self) -> float:
     """Computes the percentage better score.
@@ -1003,13 +1003,13 @@ def __call__(  # pyrefly: ignore[bad-override]
       compared_quantile: float = 0.5,
       steps_cutoff: Optional[int] = None,
   ) -> ConvergenceComparator:
-    return OptimalityGapWinRateComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=compared_curve,  # pyrefly: ignore[unexpected-keyword]
-        baseline_quantile=baseline_quantile,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=compared_quantile,  # pyrefly: ignore[unexpected-keyword]
-        name='optimality_gap_win_rate',  # pyrefly: ignore[unexpected-keyword]
-        steps_cutoff=steps_cutoff,  # pyrefly: ignore[unexpected-keyword]
+    return OptimalityGapWinRateComparator(
+        baseline_curve=baseline_curve,
+        compared_curve=compared_curve,
+        baseline_quantile=baseline_quantile,
+        compared_quantile=compared_quantile,
+        name='optimality_gap_win_rate',
+        steps_cutoff=steps_cutoff,
     )
 
 
@@ -1024,13 +1024,13 @@ def __call__(  # pyrefly: ignore[bad-override]
       compared_quantile: float = 0.5,
       steps_cutoff: Optional[int] = None,
   ) -> ConvergenceComparator:
-    return OptimalityGapGainComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=compared_curve,  # pyrefly: ignore[unexpected-keyword]
-        baseline_quantile=baseline_quantile,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=compared_quantile,  # pyrefly: ignore[unexpected-keyword]
-        name='optimality_gap_gain',  # pyrefly: ignore[unexpected-keyword]
-        steps_cutoff=steps_cutoff,  # pyrefly: ignore[unexpected-keyword]
+    return OptimalityGapGainComparator(
+        baseline_curve=baseline_curve,
+        compared_curve=compared_curve,
+        baseline_quantile=baseline_quantile,
+        compared_quantile=compared_quantile,
+        name='optimality_gap_gain',
+        steps_cutoff=steps_cutoff,
     )
 
 
@@ -1048,14 +1048,14 @@ def __call__(  # pyrefly: ignore[bad-override]
       compared_quantile: float = 0.5,
       steps_cutoff: Optional[int] = None,
   ) -> ConvergenceComparator:
-    return WinRateConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=compared_curve,  # pyrefly: ignore[unexpected-keyword]
-        baseline_quantile=baseline_quantile,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=compared_quantile,  # pyrefly: ignore[unexpected-keyword]
-        name='convergence_curve_win_rate',  # pyrefly: ignore[unexpected-keyword]
+    return WinRateConvergenceCurveComparator(
+        baseline_curve=baseline_curve,
+        compared_curve=compared_curve,
+        baseline_quantile=baseline_quantile,
+        compared_quantile=compared_quantile,
+        name='convergence_curve_win_rate',
         comparison_mode=self.comparison_mode,
-        steps_cutoff=steps_cutoff,  # pyrefly: ignore[unexpected-keyword]
+        steps_cutoff=steps_cutoff,
     )
 
 
@@ -1072,13 +1072,13 @@ def __call__(  # pyrefly: ignore[bad-override]
       compared_quantile: float = 0.5,
       steps_cutoff: Optional[int] = None,
   ) -> ConvergenceComparator:
-    return LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=compared_curve,  # pyrefly: ignore[unexpected-keyword]
-        baseline_quantile=baseline_quantile,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=compared_quantile,  # pyrefly: ignore[unexpected-keyword]
-        name='log_eff',  # pyrefly: ignore[unexpected-keyword]
-        steps_cutoff=steps_cutoff,  # pyrefly: ignore[unexpected-keyword]
+    return LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=baseline_curve,
+        compared_curve=compared_curve,
+        baseline_quantile=baseline_quantile,
+        compared_quantile=compared_quantile,
+        name='log_eff',
+        steps_cutoff=steps_cutoff,
     )
 
 
@@ -1095,13 +1095,13 @@ def __call__(  # pyrefly: ignore[bad-override]
       compared_quantile: float = 0.5,
       steps_cutoff: Optional[int] = None,
   ) -> ConvergenceComparator:
-    return PercentageBetterConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=compared_curve,  # pyrefly: ignore[unexpected-keyword]
-        baseline_quantile=baseline_quantile,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=compared_quantile,  # pyrefly: ignore[unexpected-keyword]
-        name='pct_better',  # pyrefly: ignore[unexpected-keyword]
-        steps_cutoff=steps_cutoff,  # pyrefly: ignore[unexpected-keyword]
+    return PercentageBetterConvergenceCurveComparator(
+        baseline_curve=baseline_curve,
+        compared_curve=compared_curve,
+        baseline_quantile=baseline_quantile,
+        compared_quantile=compared_quantile,
+        name='pct_better',
+        steps_cutoff=steps_cutoff,
     )
 
 
diff --git a/vizier/_src/benchmarks/analyzers/convergence_curve_test.py b/vizier/_src/benchmarks/analyzers/convergence_curve_test.py
index d286c0b72..6745a0582 100644
--- a/vizier/_src/benchmarks/analyzers/convergence_curve_test.py
+++ b/vizier/_src/benchmarks/analyzers/convergence_curve_test.py
@@ -699,16 +699,16 @@ def setUp(self):
 
   def test_higher_quantile_curve(self):
     baseline_length = len(self._baseline_curve.xs)
-    median_score = convergence.WinRateConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve, compared_curve=self._better_curves  # pyrefly: ignore[unexpected-keyword]
+    median_score = convergence.WinRateConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve, compared_curve=self._better_curves
     ).curve()
-    reverse_median_score = convergence.WinRateConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve, compared_curve=self._worse_curves  # pyrefly: ignore[unexpected-keyword]
+    reverse_median_score = convergence.WinRateConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve, compared_curve=self._worse_curves
     ).curve()
-    higher_quantile_score = convergence.WinRateConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=self._better_curves,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=0.9,  # pyrefly: ignore[unexpected-keyword]
+    higher_quantile_score = convergence.WinRateConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve,
+        compared_curve=self._better_curves,
+        compared_quantile=0.9,
     ).curve()
 
     self.assertEqual(median_score.ys.shape, (1, baseline_length))
@@ -722,22 +722,22 @@ def test_higher_quantile_curve(self):
   def test_get_winrate_score(self):
     # Higher compared quantile should increase score. Higher baseline
     # quantile should decrease score.
-    median_score = convergence.WinRateConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve, compared_curve=self._better_curves  # pyrefly: ignore[unexpected-keyword]
+    median_score = convergence.WinRateConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve, compared_curve=self._better_curves
     ).score()
     self.assertGreaterEqual(
-        convergence.WinRateConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-            baseline_curve=self._baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-            compared_curve=self._better_curves,  # pyrefly: ignore[unexpected-keyword]
-            compared_quantile=0.9,  # pyrefly: ignore[unexpected-keyword]
+        convergence.WinRateConvergenceCurveComparator(
+            baseline_curve=self._baseline_curve,
+            compared_curve=self._better_curves,
+            compared_quantile=0.9,
         ).score(),
         median_score,
     )
     self.assertLessEqual(
-        convergence.WinRateConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-            baseline_curve=self._baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-            compared_curve=self._better_curves,  # pyrefly: ignore[unexpected-keyword]
-            baseline_quantile=0.9,  # pyrefly: ignore[unexpected-keyword]
+        convergence.WinRateConvergenceCurveComparator(
+            baseline_curve=self._baseline_curve,
+            compared_curve=self._better_curves,
+            baseline_quantile=0.9,
         ).score(),
         median_score,
     )
@@ -768,13 +768,13 @@ def setUp(self):
 
   def test_get_relative_efficiency_curve(self):
     baseline_length = len(self._baseline_curve.xs)
-    rel_effiency = convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve, compared_curve=self._better_curves  # pyrefly: ignore[unexpected-keyword]
+    rel_effiency = convergence.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve, compared_curve=self._better_curves
     ).curve()
-    higher_quantile = convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=self._better_curves,  # pyrefly: ignore[unexpected-keyword]
-        compared_quantile=0.9,  # pyrefly: ignore[unexpected-keyword]
+    higher_quantile = convergence.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve,
+        compared_curve=self._better_curves,
+        compared_quantile=0.9,
     ).curve()
 
     self.assertEqual(rel_effiency.ys.shape, (1, baseline_length))
@@ -790,8 +790,8 @@ def test_get_relative_efficiency_flat(self):
         ys=np.array([4.0, 3.0, 2.0] + [1.5] * 17).reshape(1, 20),
         trend=convergence.ConvergenceCurve.YTrend.DECREASING,
     )
-    self_eff = convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=flat_curve, compared_curve=flat_curve  # pyrefly: ignore[unexpected-keyword]
+    self_eff = convergence.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=flat_curve, compared_curve=flat_curve
     ).curve()
     # Relative efficiency of a curve on itself is close to 0.
     self.assertAlmostEqual(np.linalg.norm(self_eff.ys), 0.0, delta=0.1)
@@ -805,9 +805,9 @@ def test_get_relative_efficiency_short_curve(self):
         trend=convergence.ConvergenceCurve.YTrend.DECREASING,
     )
     max_score = 10.3
-    short_efficiency = convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-        compared_curve=short_curve,  # pyrefly: ignore[unexpected-keyword]
+    short_efficiency = convergence.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve,
+        compared_curve=short_curve,
         max_score=max_score,
     ).curve()
     self.assertEqual(short_efficiency.ys.shape, (1, baseline_length))
@@ -816,32 +816,32 @@ def test_get_relative_efficiency_short_curve(self):
   def test_get_efficiency_score(self):
     # Higher compared quantile should increase score. Higher baseline
     # quantile should decrease score.
-    median_score = convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve, compared_curve=self._better_curves  # pyrefly: ignore[unexpected-keyword]
+    median_score = convergence.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve, compared_curve=self._better_curves
     ).score()
     self.assertGreater(
-        convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-            baseline_curve=self._baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-            compared_curve=self._better_curves,  # pyrefly: ignore[unexpected-keyword]
-            compared_quantile=0.9,  # pyrefly: ignore[unexpected-keyword]
+        convergence.LogEfficiencyConvergenceCurveComparator(
+            baseline_curve=self._baseline_curve,
+            compared_curve=self._better_curves,
+            compared_quantile=0.9,
         ).score(),
         median_score,
     )
     self.assertLess(
-        convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-            baseline_curve=self._baseline_curve,  # pyrefly: ignore[unexpected-keyword]
-            compared_curve=self._better_curves,  # pyrefly: ignore[unexpected-keyword]
-            baseline_quantile=0.9,  # pyrefly: ignore[unexpected-keyword]
+        convergence.LogEfficiencyConvergenceCurveComparator(
+            baseline_curve=self._baseline_curve,
+            compared_curve=self._better_curves,
+            baseline_quantile=0.9,
         ).score(),
         median_score,
     )
 
   def test_effiency_score_symmetry(self):
-    base_score = convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._baseline_curve, compared_curve=self._better_curves  # pyrefly: ignore[unexpected-keyword]
+    base_score = convergence.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=self._baseline_curve, compared_curve=self._better_curves
     ).score()
-    reversed_score = convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-        baseline_curve=self._better_curves, compared_curve=self._baseline_curve  # pyrefly: ignore[unexpected-keyword]
+    reversed_score = convergence.LogEfficiencyConvergenceCurveComparator(
+        baseline_curve=self._better_curves, compared_curve=self._baseline_curve
     ).score()
     self.assertAlmostEqual(base_score, -reversed_score, delta=0.01)
 
@@ -862,14 +862,14 @@ def test_efficiency_score_value(self):
     # Efficiency score for exponential curves can be approximated.
 
     self.assertGreater(
-        convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-            baseline_curve=self._baseline_curve, compared_curve=better_curves  # pyrefly: ignore[unexpected-keyword]
+        convergence.LogEfficiencyConvergenceCurveComparator(
+            baseline_curve=self._baseline_curve, compared_curve=better_curves
         ).score(),
         0.4,
     )
     self.assertLess(
-        convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-            baseline_curve=self._baseline_curve, compared_curve=worse_curves  # pyrefly: ignore[unexpected-keyword]
+        convergence.LogEfficiencyConvergenceCurveComparator(
+            baseline_curve=self._baseline_curve, compared_curve=worse_curves
         ).score(),
         -0.4,
     )
@@ -879,8 +879,8 @@ def test_comparator_failure(self):
         xs=self._baseline_curve.xs, ys=self._baseline_curve.ys
     )
     with self.assertRaisesRegex(ValueError, 'increasing or decreasing'):
-      convergence.LogEfficiencyConvergenceCurveComparator(  # pyrefly: ignore[missing-argument]
-          baseline_curve=unknown_curve, compared_curve=self._baseline_curve  # pyrefly: ignore[unexpected-keyword]
+      convergence.LogEfficiencyConvergenceCurveComparator(
+          baseline_curve=unknown_curve, compared_curve=self._baseline_curve
       )
 
 
diff --git a/vizier/_src/benchmarks/analyzers/state_analyzer.py b/vizier/_src/benchmarks/analyzers/state_analyzer.py
index 3bb72df15..1b40c0972 100644
--- a/vizier/_src/benchmarks/analyzers/state_analyzer.py
+++ b/vizier/_src/benchmarks/analyzers/state_analyzer.py
@@ -297,7 +297,7 @@ def add_comparison_metrics(
         analyzed_records.append(
             BenchmarkRecord(
                 algorithm=algorithm_name,
-                experimenter_metadata=vz.Metadata(json.loads(experimenter_key)),  # pytype: disable=wrong-arg-types  # pandas-drop-duplicates-overloads
+                experimenter_metadata=vz.Metadata(json.loads(experimenter_key)),  # pyrefly: ignore[bad-argument-type]
                 plot_elements=elems_dict,
             )
         )
diff --git a/vizier/_src/pyvizier/oss/automated_stopping.py b/vizier/_src/pyvizier/oss/automated_stopping.py
index 38472256a..27756ce38 100644
--- a/vizier/_src/pyvizier/oss/automated_stopping.py
+++ b/vizier/_src/pyvizier/oss/automated_stopping.py
@@ -34,12 +34,12 @@ class AutomatedStoppingConfig:
   def default_stopping_spec(cls) -> 'AutomatedStoppingConfig':
     """Use Vizier's default early stopping."""
     config = study_pb2.StudySpec.DefaultEarlyStoppingSpec()
-    return cls(proto=config)  # pyrefly: ignore[missing-argument, unexpected-keyword]
+    return cls(proto=config)
 
   @classmethod
   def from_proto(
       cls, proto: AutomatedStoppingConfigProto) -> 'AutomatedStoppingConfig':
-    return cls(proto=proto)  # pyrefly: ignore[missing-argument, unexpected-keyword]
+    return cls(proto=proto)
 
   def to_proto(self) -> AutomatedStoppingConfigProto:
     """Returns this object as a proto."""
diff --git a/vizier/_src/pyvizier/oss/proto_converters.py b/vizier/_src/pyvizier/oss/proto_converters.py
index 6458cab90..8f1c36584 100644
--- a/vizier/_src/pyvizier/oss/proto_converters.py
+++ b/vizier/_src/pyvizier/oss/proto_converters.py
@@ -649,7 +649,7 @@ def from_proto(cls, proto: study_pb2.Trial) -> trial.Trial:
         final_measurement=final_measurement,  # pyrefly: ignore[unexpected-keyword]
         measurements=measurements,  # pyrefly: ignore[unexpected-keyword]
         metadata=metadata,  # pyrefly: ignore[unexpected-keyword]
-    )  # pytype: disable=wrong-arg-types
+    )
 
   @classmethod
   def from_protos(cls, protos: Iterable[study_pb2.Trial]) -> List[trial.Trial]:
diff --git a/vizier/_src/pyvizier/oss/study_config.py b/vizier/_src/pyvizier/oss/study_config.py
index ad7f9af0e..81af59e1e 100644
--- a/vizier/_src/pyvizier/oss/study_config.py
+++ b/vizier/_src/pyvizier/oss/study_config.py
@@ -314,7 +314,7 @@ def _trial_to_external_values(
       if pc.external_type is None:
         external_value = remaining_parameters[pc.name].value
       else:
-        external_value = remaining_parameters[pc.name].cast(pc.external_type)  # pytype: disable=wrong-arg-types
+        external_value = remaining_parameters[pc.name].cast(pc.external_type)
       external_values[pc.name] = external_value
       remaining_parameters.pop(pc.name)
     return external_values
diff --git a/vizier/_src/pyvizier/shared/common.py b/vizier/_src/pyvizier/shared/common.py
index f0818ef39..693b86a25 100644
--- a/vizier/_src/pyvizier/shared/common.py
+++ b/vizier/_src/pyvizier/shared/common.py
@@ -366,7 +366,7 @@ def get_proto(self, key: str, *, cls: Type[_M]) -> Optional[_M]:
     if isinstance(value, cls):
       # Starting from 3.10, pytype supports typeguard, which obsoletes
       # the need for the `pytype:disable` clause.
-      return value  # pytype: disable=bad-return-type
+      return value
     if isinstance(value, any_pb2.Any):
       # `value` is an Any proto potentially packing `cls`.
       message = cls()
@@ -416,7 +416,7 @@ def get_or_error(self, key: str, *, cls: Type[T] = str) -> T:
     if isinstance(value, cls):
       # Starting from 3.10, pytype supports typeguard, which obsoletes
       # the need for the `pytype:disable` clause.
-      return value  # pytype: disable=bad-return-type
+      return value
     elif isinstance(value, any_pb2.Any):
       # `value` is an Any proto potentially packing `cls`.
       message = cls()
diff --git a/vizier/_src/pyvizier/shared/common_test.py b/vizier/_src/pyvizier/shared/common_test.py
index bd211b978..10dda6d73 100644
--- a/vizier/_src/pyvizier/shared/common_test.py
+++ b/vizier/_src/pyvizier/shared/common_test.py
@@ -284,7 +284,7 @@ def test_assign_proto(self):
       self.assertEqual(dur_out.seconds, 60)
 
   def test_get_or_error(self):
-    m0 = common.Metadata({'foo': 33, 'bar': 'Z'})  # pyrefly: ignore[bad-argument-type, bad-assignment]
+    m0 = common.Metadata({'foo': 33, 'bar': 'Z'})  # pyrefly: ignore[bad-assignment]
     m0['gleep'] = duration_pb2.Duration(seconds=60)
     self.assertEqual(m0.get_or_error('foo', cls=int), 33)
     self.assertEqual(m0.get_or_error('bar'), 'Z')
diff --git a/vizier/_src/pyvizier/shared/context.py b/vizier/_src/pyvizier/shared/context.py
index 308bea40e..73df3fcfd 100644
--- a/vizier/_src/pyvizier/shared/context.py
+++ b/vizier/_src/pyvizier/shared/context.py
@@ -42,8 +42,10 @@ class Context:
       validator=attr.validators.deep_mapping(
           key_validator=attr.validators.instance_of(str),
           value_validator=attr.validators.instance_of(ParameterValue),
-          mapping_validator=attr.validators.instance_of(dict)),
-      on_setattr=attr.setters.validate)  # pytype: disable=wrong-arg-types
+          mapping_validator=attr.validators.instance_of(dict),
+      ),
+      on_setattr=attr.setters.validate,
+  )
 
   metadata: Metadata = attr.ib(
       init=True,
@@ -59,5 +61,7 @@ class Context:
       validator=attr.validators.deep_mapping(
           key_validator=attr.validators.instance_of(str),
           value_validator=attr.validators.instance_of(str),
-          mapping_validator=attr.validators.instance_of(dict)),
-      on_setattr=attr.setters.validate)  # pytype: disable=wrong-arg-types
+          mapping_validator=attr.validators.instance_of(dict),
+      ),
+      on_setattr=attr.setters.validate,
+  )
diff --git a/vizier/_src/pyvizier/shared/parameter_config.py b/vizier/_src/pyvizier/shared/parameter_config.py
index fa73d29bf..30797feb3 100644
--- a/vizier/_src/pyvizier/shared/parameter_config.py
+++ b/vizier/_src/pyvizier/shared/parameter_config.py
@@ -344,15 +344,15 @@ def factory(
     if default_value is not None:
       default_value = _get_default_value(inferred_type, default_value)
 
-    pc = cls(  # pyrefly: ignore[missing-argument]
-        name=name,  # pyrefly: ignore[unexpected-keyword]
-        type=inferred_type,  # pyrefly: ignore[unexpected-keyword]
-        bounds=bounds,  # pyrefly: ignore[unexpected-keyword]
-        feasible_values=feasible_values,  # pyrefly: ignore[bad-argument-type, unexpected-keyword]
-        scale_type=scale_type,  # pyrefly: ignore[unexpected-keyword]
-        default_value=default_value,  # pyrefly: ignore[unexpected-keyword]
+    pc = cls(
+        name=name,
+        type=inferred_type,
+        bounds=bounds,
+        feasible_values=feasible_values,  # pyrefly: ignore[bad-argument-type]
+        scale_type=scale_type,
+        default_value=default_value,
         fidelity_config=fidelity_config,
-        external_type=external_type,  # pyrefly: ignore[unexpected-keyword]
+        external_type=external_type,
     )
     if children:
       pc = pc._add_children(children)
diff --git a/vizier/_src/pyvizier/shared/trial.py b/vizier/_src/pyvizier/shared/trial.py
index da3a0e4f2..772d2c911 100644
--- a/vizier/_src/pyvizier/shared/trial.py
+++ b/vizier/_src/pyvizier/shared/trial.py
@@ -99,7 +99,7 @@ def _std_not_negative(self, _, stddev: Optional[float]) -> bool:  # pyrefly: ign
     if (stddev is not None) and (not stddev >= 0):
       raise ValueError(
           'Standard deviation must be a non-negative finite number.'
-      )  # pytype: disable=bad-return-type
+      )
 
   value: float = attr.ib(
       converter=float,
@@ -147,11 +147,11 @@ def cast_as_internal(self,
     internal_type.assert_correct_type(self.value)
 
     if internal_type in (ParameterType.DOUBLE, ParameterType.DISCRETE):
-      return self.as_float  # pytype: disable=bad-return-type
+      return self.as_float  # pyrefly: ignore[bad-return]
     elif internal_type == ParameterType.INTEGER:
-      return self.as_int  # pytype: disable=bad-return-type
+      return self.as_int  # pyrefly: ignore[bad-return]
     elif internal_type == ParameterType.CATEGORICAL:
-      return self.as_str  # pytype: disable=bad-return-type
+      return self.as_str  # pyrefly: ignore[bad-return]
     else:
       raise RuntimeError(f'Unknown type {internal_type}')
 
@@ -176,11 +176,11 @@ def cast(
     if external_type == ExternalType.INTERNAL:
       return self.value
     elif external_type == ExternalType.BOOLEAN:
-      return self.as_bool  # pytype: disable=bad-return-type
+      return self.as_bool  # pyrefly: ignore[bad-return]
     elif external_type == ExternalType.INTEGER:
-      return self.as_int  # pytype: disable=bad-return-type
+      return self.as_int  # pyrefly: ignore[bad-return]
     elif external_type == ExternalType.FLOAT:
-      return self.as_float  # pytype: disable=bad-return-type
+      return self.as_float  # pyrefly: ignore[bad-return]
     else:
       raise ValueError(
           'Unknown external type enum value: {}.'.format(external_type))
@@ -362,7 +362,7 @@ class ParameterDict(abc.MutableMapping):
 
   def as_dict(self) -> Dict[str, ParameterValueTypes]:
     """Returns the dict of parameter names to raw values."""
-    return {k: self.get_value(k) for k in self._items}  # pytype: disable=bad-return-type
+    return {k: self.get_value(k) for k in self._items}  # pyrefly: ignore[bad-return]
 
   def __init__(self, iterable: Any = tuple(), **kwargs):
     self.__attrs_init__()  # pyrefly: ignore[missing-attribute]
@@ -411,7 +411,8 @@ class TrialSuggestion:
       init=True,
       factory=ParameterDict,
       converter=ParameterDict,
-      validator=attr.validators.instance_of(ParameterDict))  # pytype: disable=wrong-arg-types
+      validator=attr.validators.instance_of(ParameterDict),
+  )
 
   metadata: common.Metadata = attr.field(
       init=True,
@@ -486,8 +487,9 @@ class Trial(TrialSuggestion):
       validator=attr.validators.deep_mapping(
           key_validator=attr.validators.instance_of(str),
           value_validator=attr.validators.instance_of(str),
-          mapping_validator=attr.validators.instance_of(dict)),
-  )  # pytype: disable=wrong-arg-types
+          mapping_validator=attr.validators.instance_of(dict),
+      ),
+  )
 
   final_measurement: Optional[Measurement] = attr.ib(
       init=True,
diff --git a/vizier/pyvizier/converters/core_test.py b/vizier/pyvizier/converters/core_test.py
index cacad0011..09eca2bd5 100644
--- a/vizier/pyvizier/converters/core_test.py
+++ b/vizier/pyvizier/converters/core_test.py
@@ -140,7 +140,7 @@ def test_multi_metrics(self):
     problem = pyvizier.ProblemStatement(
         search_space=search_space,
         metric_information=pyvizier.MetricsConfig(
-            metrics=[  # pyrefly: ignore[unexpected-keyword]
+            metrics=[
                 pyvizier.MetricInformation(
                     'obj1', goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE
                 ),
diff --git a/vizier/pyvizier/converters/jnp_converters.py b/vizier/pyvizier/converters/jnp_converters.py
index 15c966ee5..29e83bed4 100644
--- a/vizier/pyvizier/converters/jnp_converters.py
+++ b/vizier/pyvizier/converters/jnp_converters.py
@@ -41,8 +41,8 @@ def __init__(
       experimental_override: str = '',
       *,
       padding_schedule: padding.PaddingSchedule = padding.PaddingSchedule(
-          num_trials=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
-          num_features=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
+          num_trials=padding.PaddingType.NONE,
+          num_features=padding.PaddingType.NONE,
       ),
   ):
     """SHOULD NOT BE USED! Use factory classmethods e.g. from_study_config."""
@@ -89,8 +89,8 @@ def from_study_config(
       *,
       scale: bool = True,
       padding_schedule: padding.PaddingSchedule = padding.PaddingSchedule(
-          num_trials=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
-          num_features=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
+          num_trials=padding.PaddingType.NONE,
+          num_features=padding.PaddingType.NONE,
       ),
       pad_oovs: bool = True,
       max_discrete_indices: int = 0,
@@ -201,10 +201,10 @@ def create_output_converter(metric):
         [create_input_converter(p) for p in sc.search_space.parameters],
         [create_output_converter(m) for m in sc.metric_information],
     )
-    return cls(  # pyrefly: ignore[missing-argument]
+    return cls(
         TrialToContinuousAndCategoricalConverter(converter),
-        problem=problem,  # pyrefly: ignore[unexpected-keyword]
-        padding_schedule=padding_schedule,  # pyrefly: ignore[unexpected-keyword]
+        problem=problem,
+        padding_schedule=padding_schedule,
     )
 
   def to_features(self, trials: Sequence[vz.TrialSuggestion]) -> vt.ModelInput:
diff --git a/vizier/pyvizier/converters/jnp_converters_test.py b/vizier/pyvizier/converters/jnp_converters_test.py
index a0328f2de..5baca4d68 100644
--- a/vizier/pyvizier/converters/jnp_converters_test.py
+++ b/vizier/pyvizier/converters/jnp_converters_test.py
@@ -38,9 +38,9 @@ def test_padding(self):
     """Tests various padding schedules."""
 
     padding_schedule = padding.PaddingSchedule(
-        num_trials=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
-        num_features=padding.PaddingType.MULTIPLES_OF_10,  # pyrefly: ignore[unexpected-keyword]
-        num_metrics=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
+        num_trials=padding.PaddingType.POWERS_OF_2,
+        num_features=padding.PaddingType.MULTIPLES_OF_10,
+        num_metrics=padding.PaddingType.NONE,
     )
     problem = vz.ProblemStatement(test_studies.flat_space_with_all_types())
     problem.metric_information.append(
@@ -223,7 +223,7 @@ def test_multi_metrics(self):
     problem = vz.ProblemStatement(
         search_space=search_space,
         metric_information=vz.MetricsConfig(
-            metrics=[  # pyrefly: ignore[unexpected-keyword]
+            metrics=[
                 vz.MetricInformation(
                     'obj1', goal=vz.ObjectiveMetricGoal.MAXIMIZE
                 ),
diff --git a/vizier/pyvizier/converters/padding_test.py b/vizier/pyvizier/converters/padding_test.py
index c7b960502..7bfa12d8c 100644
--- a/vizier/pyvizier/converters/padding_test.py
+++ b/vizier/pyvizier/converters/padding_test.py
@@ -30,9 +30,9 @@ def test_padding(self):
     labels = np.random.randn(13, 3)
 
     schedule = padding.PaddingSchedule(
-        num_trials=padding.PaddingType.MULTIPLES_OF_10,  # pyrefly: ignore[unexpected-keyword]
-        num_features=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
-        num_metrics=padding.PaddingType.POWERS_OF_2,  # pyrefly: ignore[unexpected-keyword]
+        num_trials=padding.PaddingType.MULTIPLES_OF_10,
+        num_features=padding.PaddingType.POWERS_OF_2,
+        num_metrics=padding.PaddingType.POWERS_OF_2,
     )
 
     padded = schedule.pad_features(features)
@@ -49,9 +49,9 @@ def test_nopadding(self):
     labels = np.random.randn(13, 3)
 
     schedule = padding.PaddingSchedule(
-        num_trials=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
-        num_features=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
-        num_metrics=padding.PaddingType.NONE,  # pyrefly: ignore[unexpected-keyword]
+        num_trials=padding.PaddingType.NONE,
+        num_features=padding.PaddingType.NONE,
+        num_metrics=padding.PaddingType.NONE,
     )
 
     padded = schedule.pad_features(features)