Replace # pytype: disable suppressions with # pyrefly: ignore

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(

learned_optimization/research/data_driven/mnist_projections.py

@@ -311,7 +311,7 @@ def project_batches(key, batches, step):
             p=self._project_prob).astype(subset.dtype)
         subset *= mask
 
-      key = jax.vmap(jax.random.fold_in, (None, 0))(rng_tasks_init, subset)  # pytype: disable=wrong-arg-types  # jax-types
+      key = jax.vmap(jax.random.fold_in, (None, 0))(rng_tasks_init, subset)
       key_inp, key_out, key_mask = jax.vmap(
           functools.partial(jax.random.split, num=3), out_axes=1)(
               key)

learned_optimization/research/distill/truncated_distill.py

@@ -90,11 +90,16 @@ def state_compare_mse_params(prev_src_state: UnrollState,
 
 @functools.partial(
     jax.jit,
-    static_argnames=("truncated_step", "src_truncated_step",
-                     "state_compare_loss", "with_summary", "unroll_length",
-                     "outer_param_noise"),
+    static_argnames=(
+        "truncated_step",
+        "src_truncated_step",
+        "state_compare_loss",
+        "with_summary",
+        "unroll_length",
+        "outer_param_noise",
+    ),
 )
-@functools.partial(summary.add_with_summary, static_argnums=(0, 1, 2, 3, 4))  # pyrefly: ignore[bad-specialization]
+@functools.partial(summary.add_with_summary, static_argnums=(0, 1, 2, 3, 4))
 @functools.partial(jax.value_and_grad, has_aux=True, argnums=5)
 def distill_truncated_unroll(
     truncated_step: truncated_step_mod.VectorizedTruncatedStep,

learned_optimization/research/general_lopt/hyper_v2.py

@@ -493,7 +493,7 @@ def _ff_mod(self, global_feat, extra_step_mult, p, g, m, rms, fac_g,
         for oi, oo in enumerate(o):
           o_tmp = o_tmp + oo @ w[oi]
       else:
-        o_tmp = o @ w  # pytype: disable=unsupported-operands,wrong-arg-types  # jax-operator-types
+        o_tmp = o @ w
 
       o = o_tmp + jnp.broadcast_to(b,
                                    list(o_tmp.shape[0:-1]) + [o_tmp.shape[-1]])

learned_optimization/research/hysteresis/truncated_es_shared_noise.py

@@ -97,7 +97,7 @@ def init_worker_state(self, worker_weights: gradient_learner.WorkerWeights,
         epsilons=epsilons)
 
   @profile.wrap()
-  def compute_gradient_estimate(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def compute_gradient_estimate(
       self,
       worker_weights,
       key: PRNGKey,

learned_optimization/research/jaxnerf/jaxnerf.py

@@ -123,7 +123,7 @@ def loss_with_state(self, params, state, key, data):
   def loss_with_state_and_aux(self, params, state, key, data):
     return self.loss(params, key, data), state, {}
 
-  def loss(self, params, key, data) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+  def loss(self, params, key, data) -> jnp.ndarray:  # pyrefly: ignore[bad-override]
     loss_val, _ = _nerf_loss_fn(self.model, key, params, data,
                                 self.cfg.randomized, self.cfg.weight_decay_mult)
     return jnp.mean(loss_val)
@@ -162,7 +162,7 @@ def tree_sum_fn(fn):
       tree_sum_fn(lambda z: jnp.sum(z**2)) /
       tree_sum_fn(lambda z: jnp.prod(jnp.array(z.shape))))
 
-  stats = utils.Stats(  # pytype: disable=wrong-arg-types  # jax-types
+  stats = utils.Stats(
       loss=loss, psnr=psnr, loss_c=loss_c, psnr_c=psnr_c, weight_l2=weight_l2
   )
   return loss + loss_c + weight_decay_mult * weight_l2, stats

learned_optimization/tasks/base.py

@@ -44,13 +44,13 @@ def loss_with_state(self, params: Params, state: ModelState, key: PRNGKey,
                       data: Batch) -> Tuple[jnp.ndarray, ModelState]:
     if state is not None:
       raise ValueError("Define a custom loss_with_state when using a state!")
-    return self.loss(params, key, data), None  # pytype: disable=bad-return-type  # jax-ndarray
+    return self.loss(params, key, data), None  # pyrefly: ignore[bad-return]
 
   def loss_and_aux(
       self, params: Params, key: PRNGKey,
       data: Batch) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:
     loss = self.loss(params, key, data)
-    return loss, {}  # pytype: disable=bad-return-type  # jax-ndarray
+    return loss, {}  # pyrefly: ignore[bad-return]
 
   def loss_with_state_and_aux(
       self, params: Params, state: ModelState, key: PRNGKey,

learned_optimization/tasks/datasets/base.py

@@ -460,10 +460,14 @@ def make_python_iter(split: str) -> Iterator[Batch]:
     filenames = _tfrecord_filenames_from_dataset_name(datasetname, split)
 
     filenames = [tf.convert_to_tensor(filename) for filename in filenames]
-    filenames = tf.data.Dataset.from_tensor_slices(filenames).repeat(  # pyrefly: ignore[bad-argument-type]
-        -1).shuffle(len(filenames) * 2)
-    ds = tf.data.TFRecordDataset(  # pyrefly: ignore[bad-instantiation]
-        filenames, compression_type="GZIP", num_parallel_reads=4)
+    filenames = (
+        tf.data.Dataset.from_tensor_slices(filenames)
+        .repeat(-1)
+        .shuffle(len(filenames) * 2)
+    )
+    ds = tf.data.TFRecordDataset(
+        filenames, compression_type="GZIP", num_parallel_reads=4
+    )
 
     features = {
         "image": tf.io.FixedLenFeature([], dtype=tf.string),

learned_optimization/tasks/es_wrapper.py

@@ -52,7 +52,7 @@ def _fn(key: PRNGKey) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
     pos = _sample_perturbations(theta, key, std=std)
     p_theta = jax.tree_util.tree_map(lambda t, a: t + a, theta, pos)
     n_theta = jax.tree_util.tree_map(lambda t, a: t - a, theta, pos)
-    return pos, p_theta, n_theta  # pyrefly: ignore[bad-return]
+    return pos, p_theta, n_theta
 
   keys = jax.random.split(key, num_samples)
   vec_pos, vec_p_theta, vec_n_theta = jax.vmap(_fn)(keys)

learned_optimization/tasks/fixed/image_mlp.py

@@ -57,7 +57,7 @@ def init(self, key: PRNGKey) -> Any:
                                    self.datasets.abstract_batch)  # pyrefly: ignore[missing-attribute]
     return self._mod.init(key, batch["image"])
 
-  def loss(self, params: Params, key: PRNGKey, data: Any) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+  def loss(self, params: Params, key: PRNGKey, data: Any) -> jnp.ndarray:  # pyrefly: ignore[bad-override]
     num_classes = self.datasets.extra_info["num_classes"]  # pyrefly: ignore[missing-attribute]
     logits = self._mod.apply(params, key, data["image"])
     labels = jax.nn.one_hot(data["label"], num_classes)

learned_optimization/tasks/fixed/rnn_lm.py

@@ -81,7 +81,7 @@ def init(self, key: PRNGKey) -> base.Params:
                                    self.datasets.abstract_batch)  # pyrefly: ignore[missing-attribute]
     return self._mod.init(key, batch["obs"])
 
-  def loss(self, params: Params, key: PRNGKey, data: Any) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+  def loss(self, params: Params, key: PRNGKey, data: Any) -> jnp.ndarray:  # pyrefly: ignore[bad-override]
     obs = data["obs"]
     target = data["target"]
 
@@ -138,7 +138,7 @@ def _delay(klass, *args):
 def _partial(rnn_fn, embedding_dim, vocab_size, datasets):
 
   def tmp_fn():
-    return TeacherForcedRNNLM(  # pytype: disable=wrong-arg-types
+    return TeacherForcedRNNLM(
         rnn_fn,
         embedding_dim=embedding_dim,
         vocab_size=vocab_size,

learned_optimization/tasks/quadratics.py

@@ -184,7 +184,7 @@ class _Task(base.Task):
       """Generated Task."""
       datasets = ds
 
-      def loss(self, params: Any, key: PRNGKey, data: Any) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+      def loss(self, params: Any, key: PRNGKey, data: Any) -> jnp.ndarray:  # pyrefly: ignore[bad-override]
         return jnp.sum(jnp.square(task_params - params)) + data
 
       def loss_and_aux(

learned_optimization/tasks/task_augmentation.py

@@ -90,7 +90,7 @@ def loss_with_state_and_aux(
     params = jax.tree_util.tree_map(lambda x, scale: x * scale, params, scales)
     return self.task.loss_with_state_and_aux(params, state, key, data)
 
-  def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+  def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray:  # pyrefly: ignore[bad-override]
     loss, _, _ = self.loss_with_state_and_aux(params, None, key, data)
     return loss
 
@@ -509,9 +509,9 @@ def init(self, key):
 
   def init_with_state(self, key):
     params, state = self.task.init_with_state(key)
-    buffer = self.buffer.init()  # pytype: disable=attribute-error  # jax-api-types
+    buffer = self.buffer.init()
     for _ in range(self.delay_steps):
-      buffer = self.buffer.add(buffer, params)  # pytype: disable=attribute-error  # jax-api-types
+      buffer = self.buffer.add(buffer, params)
     return params, (buffer, state)
 
   def loss(self, params, key, data):
@@ -526,10 +526,10 @@ def loss_with_aux(self, params, key, data):
 
   def loss_with_state_and_aux(self, params, state, key, data):
     buffer, state = state
-    last_entry = self.buffer.gather_from_present(buffer, -self.delay_steps + 1)  # pytype: disable=attribute-error  # jax-api-types
+    last_entry = self.buffer.gather_from_present(buffer, -self.delay_steps + 1)
     loss, next_state, aux = self.task.loss_with_state_and_aux(
         last_entry, state, key, data)
-    buffer = self.buffer.add(buffer, params)  # pytype: disable=attribute-error  # jax-api-types
+    buffer = self.buffer.add(buffer, params)
     return loss, (buffer, next_state), aux
 
 

learned_optimization/time_filter/timings.py

@@ -252,7 +252,7 @@ def timing_for_iterator(it: Iterator[Any],
       break
 
   dtimes = onp.diff(times)
-  return onp.mean(dtimes), onp.std(dtimes) / onp.sqrt(len(dtimes))  # pyrefly: ignore[bad-return]
+  return onp.mean(dtimes), onp.std(dtimes) / onp.sqrt(len(dtimes))
 
 
 def task_family_runtime_stats(
@@ -326,4 +326,4 @@ def task_family_runtime_stats(
 
   if clear_buffers:
     _jax_clear_device_buffers()
-  return ret  # pytype: disable=bad-return-type
+  return ret