Replace # pytype: disable suppressions with # pyrefly: ignore

PiperOrigin-RevId: 993589752

Author: oprypin

learned_optimization/baselines/hparam_sets.py

@@ -121,7 +121,7 @@ def RMSPropMomLR_10000_R5(task_name: str) -> HParamSet:  # pylint: disable=inval
   reps = 5
   cfgs = _lr_cfgs(task_name, "RMSProp", 10000)
   for c in cfgs:
-    c["RMSProp.momentum"] = 0.9  # pytype: disable=unsupported-operands
+    c["RMSProp.momentum"] = 0.9  # pyrefly: ignore[unsupported-operation]
   paths = [(_save_dir_from_cfg(c), reps) for c in cfgs]
   return list(cfgs) * reps, paths
 
@@ -403,7 +403,7 @@ def SM3beta2_999_LR_10000_R5(task_name: str) -> HParamSet:  # pylint: disable=in
   reps = 5
   cfgs = _lr_cfgs(task_name, "SM3", 10000, opt_name_override="SM3beta2_999")
   for c in cfgs:
-    c["SM3.b2"] = 0.999  # pytype: disable=unsupported-operands
+    c["SM3.b2"] = 0.999  # pyrefly: ignore[unsupported-operation]
   paths = [(_save_dir_from_cfg(c), reps) for c in cfgs]
   return list(cfgs) * reps, paths
 

learned_optimization/baselines/run_trainer.py

@@ -119,7 +119,8 @@ def inner_train_task(
       eval_batches=eval_batches,
       last_eval_batches=last_eval_batches,
       eval_task=eval_task,
-      device=device)  # type: MutableMapping[str, Any] # pytype: disable=annotation-type-mismatch
+      device=device,
+  )  # type: MutableMapping[str, Any]
 
   results["gin_operative_config"] = gin.operative_config_str()
 

learned_optimization/optimizers/base.py

@@ -140,7 +140,7 @@ def init(self, params, model_state=None, num_steps=None, **kwargs):  # pyrefly:
         dir_opt_state=self.direction_opt.init(
             params, model_state=model_state, num_steps=num_steps, **kwargs))
 
-  def update(self, opt_state, grad, model_state=None, **kwargs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def update(self, opt_state, grad, model_state=None, **kwargs):  # pyrefly: ignore[bad-override]
     base_params = opt_state.params
 
     next_mag_opt_state = self.magnitude_opt.update(
@@ -175,28 +175,32 @@ def update(self, opt_state, grad, model_state=None, **kwargs):  # pytype: disabl
 
 
 def SGD(*args, **kwargs):  # pylint: disable=invalid-name
-  from learned_optimization.optimizers import optax_opts  # pytype: disable=import-error # pylint: disable=g-import-not-at-top
+  from learned_optimization.optimizers import optax_opts  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
+
   warnings.warn("SGD module has been moved to optax_opts!"
                 " Calling here from base is deprecated!")
   return optax_opts.SGD(*args, **kwargs)
 
 
 def SGDM(*args, **kwargs):  # pylint: disable=invalid-name
-  from learned_optimization.optimizers import optax_opts  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
+  from learned_optimization.optimizers import optax_opts  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
+
   warnings.warn("SGDM module has been moved to optax_opts!"
                 " Calling here from base is deprecated!")
   return optax_opts.SGDM(*args, **kwargs)
 
 
 def RMSProp(*args, **kwargs):  # pylint: disable=invalid-name
-  from learned_optimization.optimizers import optax_opts  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
+  from learned_optimization.optimizers import optax_opts  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
+
   warnings.warn("RMSProp module has been moved to optax_opts!"
                 " Calling here from base is deprecated!")
   return optax_opts.RMSProp(*args, **kwargs)
 
 
 def Adam(*args, **kwargs):  # pylint: disable=invalid-name
-  from learned_optimization.optimizers import optax_opts  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
+  from learned_optimization.optimizers import optax_opts  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
+
   warnings.warn("Adammodule has been moved to optax_opts!"
                 " Calling here from base is deprecated!")
   return optax_opts.Adam(*args, **kwargs)

learned_optimization/optimizers/learning_rate_schedules.py

@@ -77,7 +77,7 @@ def __init__(self,
     self.constant_fraction = constant_fraction
     self.warmup_fraction = warmup_fraction
 
-  def __call__(self, global_step, max_steps) -> chex.Array:  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def __call__(self, global_step, max_steps) -> chex.Array:
     def fload32(x):
       """Convert input to float32."""
       return jnp.asarray(x, dtype=onp.float32)
@@ -191,4 +191,4 @@ def __init__(self, base_lr: float, decay_amount: float):
   def __call__(self,
                global_step: chex.Array,
                max_steps: Optional[chex.Array] = None) -> chex.Array:
-    return self.base_lr * (1 - self.decay_amount)**global_step  # pyrefly: ignore[bad-return]
+    return self.base_lr * (1 - self.decay_amount) ** global_step

learned_optimization/optimizers/optax_opts.py

@@ -55,7 +55,7 @@ def init(self,
            model_state: Optional[ModelState] = None,
            num_steps: Optional[int] = None,
            key: Optional[chex.PRNGKey] = None):
-    return OptaxState(  # pytype: disable=wrong-arg-types  # jax-ndarray
+    return OptaxState(
         params=params,
         optax_opt_state=self.opt.init(params),  # pyrefly: ignore[bad-argument-type]
         state=model_state,  # pyrefly: ignore[bad-argument-type]
@@ -188,7 +188,7 @@ def __init__(self,
         optax.scale_by_adam(
             b1=beta1, b2=beta2, eps=epsilon, eps_root=epsilon_root
         ),
-        optax.scale_by_schedule(piecewise_linear(times, vals=lrs)),  # pytype: disable=wrong-arg-types  # jax-arraylike
+        optax.scale_by_schedule(piecewise_linear(times, vals=lrs)),  # pyrefly: ignore[bad-argument-type]
         optax.scale(-1),
     )
     super().__init__(opt)
@@ -395,7 +395,7 @@ def init(
     should_reshape = jax.tree_util.tree_map(lambda x: len(x.shape) == 0, params)  # pylint: disable=g-explicit-length-test
     params = jax.tree_util.tree_map(_expand_scalar, params, should_reshape)
     out = super().init(params, model_state, num_steps, key)
-    return SM3OptState(  # pytype: disable=wrong-arg-types  # jax-ndarray
+    return SM3OptState(
         params=out.params,
         state=out.state,
         optax_opt_state=out.optax_opt_state,

learned_optimization/population/examples/complex_cnn/train_threads.py

@@ -159,7 +159,7 @@ def mutate_fn(
     oml_beta2 = onp.log(1 - meta_params["beta2"])
     beta2 = 1 - onp.exp(oml_beta2 + onp.random.normal() * 0.03)
 
-    return {  # pytype: disable=bad-return-type  # jax-ndarray
+    return {
         "learning_rate": onp.exp(loglr + offset),  # pyrefly: ignore[bad-assignment]
         "beta1": beta1,  # pyrefly: ignore[bad-assignment]
         "beta2": beta2,  # pyrefly: ignore[bad-assignment]
@@ -196,7 +196,7 @@ def mutate_fn(
   }
 
   initial_population = [initial_meta for _ in range(num_workers)]
-  initial_population = [mutate_fn(m) for m in initial_population]  # pytype: disable=wrong-arg-types  # jax-ndarray
+  initial_population = [mutate_fn(m) for m in initial_population]  # pyrefly: ignore[bad-argument-type]
   population = population_mod.PopulationController(initial_population, mutator)
   server = population_mod.start_courier_server("population", population)  # pylint: disable=unused-variable
 

learned_optimization/population/mutators/fixed_schedule.py

@@ -46,7 +46,7 @@ def update(self, state: Any,
     steps = cache[worker.generation_id]
 
     # grab the last checkpoint here.
-    last_checkpoint = steps.values()[-1]  # pytype: disable=unsupported-operands
+    last_checkpoint = steps.values()[-1]  # pyrefly: ignore[bad-index]
     logging.info("Active worker: %s", str(worker))
     logging.info("last checkpoint : %s", str(last_checkpoint))
 

learned_optimization/population/mutators/single_worker_explore.py

@@ -105,7 +105,7 @@ def add_worker_to_cache(from_checkpoint: population.Checkpoint,
       state["branch_checkpoint"] = steps[0]
       state["center"] = steps[0].generation_id
 
-    last_checkpoint = steps.values()[-1]  # pytype: disable=unsupported-operands
+    last_checkpoint = steps.values()[-1]  # pyrefly: ignore[bad-index]
 
     if state["phase"] == "exploit":
       # switch to center.

learned_optimization/population/mutators/winner_take_all_genetic.py

@@ -63,11 +63,11 @@ def update(
 
       if self._steps_per_exploit:
         if (
-            cache[genid].keys()[-1] - cache[genid].keys()[0]  # pytype: disable=unsupported-operands
+            cache[genid].keys()[-1] - cache[genid].keys()[0]  # pyrefly: ignore[bad-index]
         ) < self._steps_per_exploit:
           return None, current_workers
 
-        to_test = cache[genid].keys()[0] + self._steps_per_exploit  # pytype: disable=unsupported-operands
+        to_test = cache[genid].keys()[0] + self._steps_per_exploit  # pyrefly: ignore[bad-index]
         valid_values = [
             x.value
             for (s, x) in cache[genid].items()

learned_optimization/research/brax/brax_env_truncated_step.py

@@ -152,7 +152,7 @@ def do_step(unroll_state):
       action = self.policy.apply(theta, key, unroll_state.env_state.obs)
       next_env_state = self.env.step(unroll_state.env_state, action)
 
-      out = truncated_step.TruncatedUnrollOut(  # pytype: disable=wrong-arg-types  # jax-ndarray
+      out = truncated_step.TruncatedUnrollOut(
           loss=-next_env_state.reward,  # pyrefly: ignore[bad-argument-type]
           is_done=False,  # pyrefly: ignore[bad-argument-type]
           task_param=None,
@@ -165,7 +165,7 @@ def do_step(unroll_state):
       ), out
 
     def reset(unroll_state):
-      out = truncated_step.TruncatedUnrollOut(  # pytype: disable=wrong-arg-types  # jax-ndarray
+      out = truncated_step.TruncatedUnrollOut(
           loss=0.0, is_done=True, task_param=None, iteration=0, mask=False  # pyrefly: ignore[bad-argument-type]
       )
       unroll_state = BraxEnvState(