Suppress type errors for Pyre upgrade (landing group 4) (#5298)

ax/benchmark/problems/synthetic/from_botorch.py

@@ -190,6 +190,7 @@ def create_problem_from_botorch(
     if isinstance(test_problem_class, str):
         test_problem_class = getattr(synthetic, test_problem_class)
 
+    # pyrefly: ignore [not-callable]
     test_problem = test_problem_class(**test_problem_kwargs)
     is_constrained = isinstance(test_problem, ConstrainedBaseTestProblem)
 
@@ -264,6 +265,7 @@ def create_problem_from_botorch(
     else:
         worst_feasible_value = None  # Not needed for unconstrained problems
     baseline_value = (
+        # pyrefly: ignore [missing-attribute]
         BOTORCH_BASELINE_VALUES[(test_problem_class.__name__, dim)]
         if baseline_value is None
         else baseline_value

ax/core/utils.py

@@ -733,12 +733,15 @@ def _sort_trials(
                 (
                     # if the trial is running, we set end time to now for sorting ease
                     (
-                        _time_trial_completed_safe(trial=t).timestamp()
+                        _time_trial_completed_safe(
+                            trial=t  # pyrefly: ignore [bad-argument-type]
+                        ).timestamp()  # pyrefly: ignore [bad-argument-type]
                         if not trials_are_running
                         else default_time_run_started.timestamp()
                     )
                     - _time_trial_started_safe(
-                        trial=t, default_time_run_started=default_time_run_started
+                        trial=t,  # pyrefly: ignore [bad-argument-type]
+                        default_time_run_started=default_time_run_started,  # pyrefly: ignore [bad-argument-type]
                     ).timestamp()
                 )
                 // twelve_hours_in_secs

ax/generators/discrete/ashr_utils.py

@@ -52,7 +52,9 @@ def __init__(
         if prior_vars is None:
             prior_stds = prior_grid(Y=Y, Yvar=Yvar, grid_param=eb_grid_param)
             prior_vars = prior_stds**2
+        # pyrefly: ignore [bad-assignment]
         self.prior_vars: npt.NDArray = prior_vars
+        # pyrefly: ignore [bad-argument-type]
         self.ll: npt.NDArray = marginal_densities(Y=Y, Yvar=Yvar, prior_vars=prior_vars)
 
     def posterior(self, w: npt.NDArray) -> GaussianMixture:
@@ -115,15 +117,22 @@ def fit(
         k = len(self.prior_vars)  # total number of classes
         if lambdas is None:
             lambdas = np.ones(k)  # no penalty
+        # pyrefly: ignore [bad-argument-type]
         if len(lambdas) != k:
             raise ValueError(
                 "The length of the penalty sequence should be the number of "
                 "prior classes."
             )
+        # pyrefly: ignore [missing-attribute]
         lambdas = lambdas.astype(np.float64)
 
         results = fit_ashr_em(
-            ll=self.ll, lambdas=lambdas, threshold=threshold, nsteps=nsteps
+            # pyrefly: ignore [bad-argument-type]
+            ll=self.ll,
+            # pyrefly: ignore [bad-argument-type]
+            lambdas=lambdas,
+            threshold=threshold,
+            nsteps=nsteps,
         )
 
         return results

ax/generators/discrete/eb_thompson.py

@@ -61,4 +61,5 @@ def _apply_shrinkage(
         Y = npY.tolist()
         npYvar = npYsem**2
         Yvar = npYvar.tolist()
+        # pyrefly: ignore [bad-return]
         return Y, Yvar

ax/generators/torch/tests/test_kernels.py

@@ -63,9 +63,9 @@ def test_scalematern_kernel_active_dims(self) -> None:
         base_kernel = assert_is_instance(covar.base_kernel, MaternKernel)
         # active_dims lands on the inner MaternKernel, and the ScaleKernel
         # inherits it (so subsetting happens exactly once at the wrapper level).
-        # pyrefly: ignore [not-callable]
+        # pyrefly: ignore [missing-attribute, not-callable]
         self.assertEqual(covar.active_dims.tolist(), active_dims)
-        # pyrefly: ignore [not-callable]
+        # pyrefly: ignore [missing-attribute, not-callable]
         self.assertEqual(base_kernel.active_dims.tolist(), active_dims)
         self.assertEqual(base_kernel.ard_num_dims, len(active_dims))
         # The kernel only consumes the active columns: perturbing an inactive
@@ -121,11 +121,11 @@ def test_scale_rbf_linear_kernel_active_dims(self) -> None:
         rbf = assert_is_instance(scale_rbf.base_kernel, RBFKernel)
         # active_dims lands on both leaves, and the ScaleKernel inherits it
         # from the wrapped RBF kernel (so subsetting happens exactly once).
-        # pyrefly: ignore [not-callable]
+        # pyrefly: ignore [missing-attribute, not-callable]
         self.assertEqual(scale_rbf.active_dims.tolist(), active_dims)
-        # pyrefly: ignore [not-callable]
+        # pyrefly: ignore [missing-attribute, not-callable]
         self.assertEqual(rbf.active_dims.tolist(), active_dims)
-        # pyrefly: ignore [not-callable]
+        # pyrefly: ignore [missing-attribute, not-callable]
         self.assertEqual(linear.active_dims.tolist(), active_dims)
         # ard_num_dims matches the number of active dims.
         self.assertEqual(rbf.ard_num_dims, len(active_dims))
@@ -283,9 +283,9 @@ def test_default_kernel(self) -> None:
                 self.assertIsNone(botorch_kernel.active_dims)
             else:
                 self.assertEqual(
-                    # pyrefly: ignore [not-callable]
+                    # pyrefly: ignore [missing-attribute, not-callable]
                     ax_kernel.active_dims.tolist(),
-                    # pyrefly: ignore [not-callable]
+                    # pyrefly: ignore [missing-attribute, not-callable]
                     botorch_kernel.active_dims.tolist(),
                 )
 
@@ -301,5 +301,5 @@ def test_default_mle(self) -> None:
             self.assertTrue((kernel.lengthscale == sqrt(2) / 10).all())
             self.assertEqual(kernel.lengthscale.shape, torch.Size([3, 1, 2]))
             self.assertFalse(hasattr(kernel, "lengthscale_prior"))
-            # pyrefly: ignore [not-callable]
+            # pyrefly: ignore [missing-attribute, not-callable]
             self.assertEqual(kernel.active_dims.tolist(), active_dims)

ax/plot/contour.py

@@ -580,7 +580,13 @@ def interact_contour_plotly(
         xbutton_data_args = {"x": [], "y": [], "z": []}
         for yvar in param_names:
             res = relativize_data(
-                f_dict[xvar][yvar], sd_dict[xvar][yvar], rel, arm_data, metric
+                # pyrefly: ignore [bad-argument-type]
+                f_dict[xvar][yvar],
+                # pyrefly: ignore [bad-argument-type]
+                sd_dict[xvar][yvar],
+                rel,
+                arm_data,
+                metric,
             )
             f_final = res[0]
             sd_final = res[1]
@@ -624,7 +630,9 @@ def interact_contour_plotly(
             {
                 "xaxis.title": short_name(xvar),
                 "xaxis2.title": short_name(xvar),
+                # pyrefly: ignore [bad-argument-type]
                 "xaxis.range": axis_range(grid_dict[xvar], is_log_dict[xvar]),
+                # pyrefly: ignore [bad-argument-type]
                 "xaxis2.range": axis_range(grid_dict[xvar], is_log_dict[xvar]),
                 "xaxis.type": "log" if is_log_dict[xvar] else "linear",
                 "xaxis2.type": "log" if is_log_dict[xvar] else "linear",
@@ -649,10 +657,14 @@ def interact_contour_plotly(
                     {
                         "yaxis.title": short_name(y_param),
                         "yaxis.range": axis_range(
-                            grid_dict[y_param], is_log_dict[y_param]
+                            # pyrefly: ignore [bad-argument-type]
+                            grid_dict[y_param],
+                            is_log_dict[y_param],
                         ),
                         "yaxis2.range": axis_range(
-                            grid_dict[y_param], is_log_dict[y_param]
+                            # pyrefly: ignore [bad-argument-type]
+                            grid_dict[y_param],
+                            is_log_dict[y_param],
                         ),
                         "yaxis.type": "log" if is_log_dict[y_param] else "linear",
                         "yaxis2.type": "log" if is_log_dict[y_param] else "linear",
@@ -785,7 +797,9 @@ def interact_contour_plotly(
     # Initially visible yvar
     yvar = param_names[1]
 
+    # pyrefly: ignore [bad-argument-type]
     xrange = axis_range(grid_dict[xvar], is_log_dict[xvar])
+    # pyrefly: ignore [bad-argument-type]
     yrange = axis_range(grid_dict[yvar], is_log_dict[yvar])
 
     xtype = "log" if is_log_dict[xvar] else "linear"

ax/utils/stats/math_utils.py

@@ -165,6 +165,7 @@ def unrelativize(
         if bias_correction:
             m_t = m_t / (1 - (sem_c / abs_mean_c) ** 2)
 
+        # pyrefly: ignore [unsupported-operation]
         var = sems_t**2
         c = m_t / mean_c
         s_t2 = var * (mean_c**2) + 2 * c * cov_means - (c**2) * (sem_c**2)