Replace # pytype: disable suppressions with # pyrefly: ignore

learned_optimization/circular_buffer.py

@@ -62,7 +62,7 @@ def build_one(x):
                 jnp.ones([self.size], dtype=jnp.int64) * -self.size))
 
   @functools.partial(jax.jit, static_argnums=(0,))
-  def add(self, state: CircularBufferState, value: T) -> CircularBufferState:  # pytype: disable=invalid-annotation
+  def add(self, state: CircularBufferState, value: T) -> CircularBufferState:
     """Construct the initial state of the circular buffer with default value."""
     idx = state.idx % self.size
 
@@ -81,7 +81,7 @@ def _reorder(self, vals, idx):
     return jnp.roll(vals, -offset, axis=0)
 
   @functools.partial(jax.jit, static_argnums=(0,))
-  def stack_with_idx(self, state: CircularBufferState) -> Tuple[T, jnp.ndarray]:  # pytype: disable=invalid-annotation
+  def stack_with_idx(self, state: CircularBufferState) -> Tuple[T, jnp.ndarray]:
     """Return raw values with integer array containing index.
 
     Args:
@@ -97,7 +97,7 @@ def stack_with_idx(self, state: CircularBufferState) -> Tuple[T, jnp.ndarray]:
     return state.values[0], jnp.where(state.values[1] == -1, -1, candidate)
 
   @functools.partial(jax.jit, static_argnums=(0,))
-  def stack_reorder(self, state: CircularBufferState) -> Tuple[T, jnp.ndarray]:  # pytype: disable=invalid-annotation
+  def stack_reorder(self, state: CircularBufferState) -> Tuple[T, jnp.ndarray]:
     """Reorder the values, and return with a mask."""
     candidate = jnp.clip((state.values[1] - state.idx + self.size), -1,
                          self.size)
@@ -107,7 +107,8 @@ def stack_reorder(self, state: CircularBufferState) -> Tuple[T, jnp.ndarray]:  #
 
   @functools.partial(jax.jit, static_argnums=(0,))
   def gather_from_present(
-      self, state: CircularBufferState, idxs: jnp.ndarray) -> T:  # pytype: disable=invalid-annotation
+      self, state: CircularBufferState, idxs: jnp.ndarray
+  ) -> T:
     """Get the values from for each idx in the past."""
     offset = (idxs % self.size)
     idx = (state.idx + offset) % self.size

learned_optimization/distributed.py

@@ -128,8 +128,10 @@ def start_server(self):
 
   def _is_step_valid(self, step: int) -> bool:
     step = onp.asarray(step)  # pyrefly: ignore[bad-assignment]
-    return (self._current_iteration >= step and  # pytype: disable=bad-return-type  # typed-numpy
-            (self._current_iteration - step) <= self._staleness)
+    return (
+        self._current_iteration >= step
+        and (self._current_iteration - step) <= self._staleness
+    )
 
   def put_grads(self, worker_id: Any, step: int, value: T):
     """Put computed gradients into learner."""

learned_optimization/eval_training.py

@@ -243,14 +243,14 @@ def single_task_training_curves(
     ret["eval/last_eval_batches"] = onp.asarray(last_eval_batches)
     ret["eval/eval_batches"] = onp.asarray(eval_batches)
     ret = {**ret, **stacked_metrics}
-  return ret  # pytype: disable=bad-return-type  # jax-ndarray
+  return ret  # pyrefly: ignore[bad-return]
 
 
 @functools.partial(jax.pmap, static_broadcasted_argnums=(1,))
 @functools.partial(jax.vmap, in_axes=(0, None))
 def _pmap_vector_random_split(key: PRNGKey, n_split: int) -> PRNGKey:
   key1, key2 = jax.random.split(key)
-  return jax.random.split(key1, n_split), key2  # pytype: disable=bad-return-type  # jax-ndarray
+  return jax.random.split(key1, n_split), key2  # pyrefly: ignore[bad-return]
 
 
 @dataclasses.dataclass

learned_optimization/filesystem.py

@@ -80,4 +80,4 @@ def remove(path: str) -> bool:
   """Remove a file."""
   if _path_on_gcp(path):
     return tf.io.gfile.remove(path)
-  return shutil.rmtree(path)  # pytype: disable=bad-return-type
+  return shutil.rmtree(path)  # pyrefly: ignore[bad-return]

learned_optimization/jax_utils.py

@@ -27,7 +27,7 @@
   from jax.extend.core import unsafe_am_i_under_a_jit_DO_NOT_USE  # pylint: disable=g-import-not-at-top
 except ImportError:
   # JAX v0.9.2 or older
-  from jax.core import unsafe_am_i_under_a_jit_DO_NOT_USE  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+  from jax.core import unsafe_am_i_under_a_jit_DO_NOT_USE  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
 
 
 def maybe_static_cond(pred, true_fn, false_fn, val):

learned_optimization/learned_optimizers/adafac_mlp_lopt.py

@@ -241,7 +241,7 @@ def _mod(self, global_feat, p, g, m, rms, fac_g, fac_vec_col, fac_vec_row,
       net = inp_stack
       for wi, (w, b) in enumerate(zip(weights, biases)):
         o_tmp = net @ w
-        net = o_tmp + jnp.broadcast_to(b, list(net.shape[0:-1]) + [w.shape[-1]])  # pytype: disable=attribute-error
+        net = o_tmp + jnp.broadcast_to(b, list(net.shape[0:-1]) + [w.shape[-1]])
 
         if wi != len(weights) - 1:
           net = jax.nn.relu(net)
@@ -270,7 +270,7 @@ def _mod(self, global_feat, p, g, m, rms, fac_g, fac_vec_col, fac_vec_row,
             if type(w) == list:  # pylint: disable=unidiomatic-typecheck
               outs.append(v * w[vi][oi])
             else:
-              outs.append(v * w[vi, oi])  # pytype: disable=unsupported-operands
+              outs.append(v * w[vi, oi])
 
           if wi == 0:
             training_step_feature = global_feat["training_step_feature"]
@@ -279,7 +279,7 @@ def _mod(self, global_feat, p, g, m, rms, fac_g, fac_vec_col, fac_vec_row,
               if type(w) == list:  # pylint: disable=unidiomatic-typecheck
                 outs.append(training_step_feature[i] * w[vi][oi])
               else:
-                outs.append(training_step_feature[i] * w[vi, oi])  # pytype: disable=unsupported-operands
+                outs.append(training_step_feature[i] * w[vi, oi])
 
           grids.append(outs)
 
@@ -405,7 +405,7 @@ def init(
             num_steps=jnp.asarray(num_steps))
 
       def update(  # pyrefly: ignore[bad-override]
-          self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+          self,
           opt_state: AdafacMLPLOptState,
           grad: opt_base.Gradient,
           loss: jnp.ndarray,

learned_optimization/learned_optimizers/adafac_nominal.py

@@ -264,7 +264,7 @@ def _mod(self, global_feat, p, g, m, rms, fac_g, fac_vec_col, fac_vec_row,
       net = inp_stack
       for wi, (w, b) in enumerate(zip(weights, biases)):
         o_tmp = net @ w
-        net = o_tmp + jnp.broadcast_to(b, list(net.shape[0:-1]) + [w.shape[-1]])  # pytype: disable=attribute-error
+        net = o_tmp + jnp.broadcast_to(b, list(net.shape[0:-1]) + [w.shape[-1]])
 
         if wi != len(weights) - 1:
           net = jax.nn.relu(net)
@@ -302,7 +302,7 @@ def _mod(self, global_feat, p, g, m, rms, fac_g, fac_vec_col, fac_vec_row,
             if type(w) == list:  # pylint: disable=unidiomatic-typecheck
               outs.append(v * w[vi][oi])
             else:
-              outs.append(v * w[vi, oi])  # pytype: disable=unsupported-operands
+              outs.append(v * w[vi, oi])
 
           if wi == 0:
             training_step_feature = global_feat["training_step_feature"]
@@ -311,7 +311,7 @@ def _mod(self, global_feat, p, g, m, rms, fac_g, fac_vec_col, fac_vec_row,
               if type(w) == list:  # pylint: disable=unidiomatic-typecheck
                 outs.append(training_step_feature[i] * w[vi][oi])
               else:
-                outs.append(training_step_feature[i] * w[vi, oi])  # pytype: disable=unsupported-operands
+                outs.append(training_step_feature[i] * w[vi, oi])
 
           grids.append(outs)
 
@@ -478,7 +478,7 @@ def init(
             num_steps=jnp.asarray(num_steps))
 
       def update(  # pyrefly: ignore[bad-override]
-          self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+          self,
           opt_state: AdafacMLPLOptState,
           grad: opt_base.Gradient,
           loss: jnp.ndarray,

learned_optimization/learned_optimizers/base.py

@@ -214,15 +214,15 @@ def init(self, params, model_state=None, num_steps=None, **kwargs):  # pyrefly:
         opt.init(params, model_state, num_steps=num_steps, **kwargs)
         for opt in self.opts
     ])
-    return SumOptimizerState(0, params, model_state, opt_states)  # pytype: disable=wrong-arg-types  # jax-ndarray
+    return SumOptimizerState(0, params, model_state, opt_states)  # pyrefly: ignore[bad-argument-type]
 
   def get_params(self, state):
     return self.opts[0].get_params(state.inner_opt_states[0])
 
   def get_state(self, state):
     return self.opts[0].get_state(state.inner_opt_states[0])
 
-  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]
     # apply to both opts
     new_opt_states = [
         opt.update(os, grad, model_state=model_state, **kwargs)

learned_optimization/learned_optimizers/common.py

@@ -66,11 +66,11 @@ def _vmap_accumulator(accumulator: Callable[[float], _InitUpdate],
   """Helper function that vmaps an accumulator fn to run on multiple decays."""
 
   def init_fn(p):
-    return jax.vmap(lambda d: accumulator(d).init(p), out_axes=-1)(decays)  # pytype: disable=wrong-arg-types  # jax-ndarray
+    return jax.vmap(lambda d: accumulator(d).init(p), out_axes=-1)(decays)
 
   def update(state, grads):
     return jax.vmap(
-        lambda s, d: accumulator(d).update(s, grads),  # pytype: disable=wrong-arg-types  # jax-ndarray
+        lambda s, d: accumulator(d).update(s, grads),
         in_axes=-1,
         out_axes=-1,
     )(state, decays)

learned_optimization/learned_optimizers/mlp_lopt.py

@@ -116,7 +116,7 @@ def init(self,
             iteration=jnp.asarray(0, dtype=jnp.int32))
 
       def update(  # pyrefly: ignore[bad-override]
-          self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+          self,
           opt_state: MLPLOptState,
           grad: Any,
           loss: float,

learned_optimization/learned_optimizers/nn_adam.py

@@ -222,22 +222,17 @@ def init(self, key: PRNGKey) -> lopt_base.MetaParams:
         lambda: self.lstm_fn().initial_state(1))[1](None, key1)
 
     return flax.core.FrozenDict({
-        "lstm_init_state":
-            lstm_inital_state,
-        "rnn_params":
-            self.rnn_init(key2, jnp.zeros([1, self.rnn_input_features]),
-                          lstm_inital_state),
-        "rnn_to_controls_params":
-            self.rnn_to_controls.init(key3,
-                                      jnp.zeros([0, self.lstm_hidden_size])),
-        "per_layer_lr":
-            _scaled_lr.forward(self.initial_learning_rate),  # pyrefly: ignore[bad-argument-type]
-        "per_layer_beta1":
-            _scaled_one_minus_log.forward(self.initial_beta1),  # pyrefly: ignore[bad-argument-type]
-        "per_layer_beta2":
-            _scaled_one_minus_log.forward(self.initial_beta2),  # pyrefly: ignore[bad-argument-type]
-        "per_layer_epsilon":
-            _scaled_epsilon.forward(self.initial_epsilon),  # pyrefly: ignore[bad-argument-type]
+        "lstm_init_state": lstm_inital_state,
+        "rnn_params": self.rnn_init(
+            key2, jnp.zeros([1, self.rnn_input_features]), lstm_inital_state
+        ),
+        "rnn_to_controls_params": self.rnn_to_controls.init(
+            key3, jnp.zeros([0, self.lstm_hidden_size])
+        ),
+        "per_layer_lr": _scaled_lr.forward(self.initial_learning_rate),
+        "per_layer_beta1": _scaled_one_minus_log.forward(self.initial_beta1),
+        "per_layer_beta2": _scaled_one_minus_log.forward(self.initial_beta2),
+        "per_layer_epsilon": _scaled_epsilon.forward(self.initial_epsilon),
     })
 
   def opt_fn(self,
@@ -287,7 +282,7 @@ def init(
             per_layer_epsilon=jax.tree_util.tree_map(
                 lambda x: theta["per_layer_epsilon"], params
             ),
-        )  # pytype: disable=wrong-arg-types
+        )
 
       def lstm_features_for_tensor(self, p: jnp.ndarray, g: jnp.ndarray,
                                    m: jnp.ndarray, rms: jnp.ndarray,

learned_optimization/learned_optimizers/opt_from_checkpoint.py

@@ -66,15 +66,15 @@ def _fn(*args, **kwargs):
   @functools.lru_cache(None)
   def __getattr__(self, *args):
     if self.has_locked:
-      ret = self.ginned.__getattribute__(*args)  # pytype: disable=attribute-error
+      ret = self.ginned.__getattribute__(*args)
     else:
       if not self.has_locked:
-        ret = self.ginned.__getattribute__(*args)  # pytype: disable=attribute-error
+        ret = self.ginned.__getattribute__(*args)
       else:
         self.has_locked = True
         with gin.config_scope(None):
           with gin.config_scope(self.scope):
-            ret = self.ginned.__getattribute__(*args)  # pytype: disable=attribute-error
+            ret = self.ginned.__getattribute__(*args)
         self.has_locked = False
     if callable(ret):
       ret = self.__wrap(ret, self.scope)

learned_optimization/learned_optimizers/rnn_mlp_lopt.py

@@ -558,7 +558,7 @@ def to_map_get_mlp_features(m, rms, g, v, ff_inputs):
               rms,
               g,
               v,
-              ff_inputs,  # pytype: disable=wrong-arg-types  # jax-ndarray
+              ff_inputs,
               opt_state.iteration,
               num_tensors,  # pyrefly: ignore[bad-argument-type]
           )

learned_optimization/outer_train.py

@@ -238,7 +238,7 @@ def metrics_and_info_from_gradients(
   max_stale = current_step - onp.min(steps)
   metrics["max_staleness"] = max_stale
 
-  return metrics, worker_ids, applied_inner_steps  # pytype: disable=bad-return-type
+  return metrics, worker_ids, applied_inner_steps  # pyrefly: ignore[bad-return]
 
 
 def maybe_resample_gradient_estimators(
@@ -406,7 +406,7 @@ def build_static_and_init_unroll_state(
     total_inner_steps = onp.asarray(0, dtype=onp.int64)
 
     with profile.Profile("grads_to_onp"):
-      to_put_grads = GradientsFromWorker(  # pytype: disable=wrong-arg-types  # jax-ndarray
+      to_put_grads = GradientsFromWorker(
           metrics=gradient_worker_out.metrics,
           worker_id=worker_id,  # pyrefly: ignore[bad-argument-type]
           total_inner_steps=total_inner_steps,  # pyrefly: ignore[bad-argument-type]
@@ -998,7 +998,7 @@ def build_static_and_init_unroll_state(worker_weights, key):
     metrics = gradient_worker_out.metrics
 
     with profile.Profile("grads_to_onp"):
-      to_put_grads = GradientsFromWorker(  # pytype: disable=wrong-arg-types  # jax-ndarray
+      to_put_grads = GradientsFromWorker(
           metrics=gradient_worker_out.metrics,
           worker_id=0,
           total_inner_steps=total_inner_steps,  # pyrefly: ignore[bad-argument-type]
@@ -1043,7 +1043,7 @@ def build_static_and_init_unroll_state(worker_weights, key):
 
     delta_time = time.time() - learner_time
     learner_time = time.time()
-    to_write = summarize_learner(  # pytype: disable=wrong-arg-types  # typed-numpy
+    to_write = summarize_learner(
         step=step,
         metrics=metrics,
         worker_ids=worker_ids,

learned_optimization/outer_trainers/common.py

@@ -85,10 +85,15 @@ def _stack(a, b, axis=0):
 
 @functools.partial(
     jax.jit,
-    static_argnames=("truncated_step", "with_summary", "unroll_length",
-                     "theta_is_vector", "wrap_step_fn"),
+    static_argnames=(
+        "truncated_step",
+        "with_summary",
+        "unroll_length",
+        "theta_is_vector",
+        "wrap_step_fn",
+    ),
 )
-@functools.partial(summary.add_with_summary, static_argnums=(0, 1, 2, 3, 9))  # pyrefly: ignore[bad-specialization]
+@functools.partial(summary.add_with_summary, static_argnums=(0, 1, 2, 3, 9))
 def truncated_unroll(
     truncated_step: truncated_step_mod.VectorizedTruncatedStep,
     unroll_length: int,

learned_optimization/outer_trainers/full_es.py

@@ -308,7 +308,7 @@ def single_vec_batch(theta, state, key_data):
 
   es_grad = jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), vec_es_grad)
 
-  return jnp.mean((pos_loss + neg_loss) / 2.0), es_grad  # pytype: disable=bad-return-type
+  return jnp.mean((pos_loss + neg_loss) / 2.0), es_grad  # pyrefly: ignore[bad-return]
 
 
 @gin.configurable

learned_optimization/outer_trainers/full_grad.py

@@ -148,7 +148,7 @@ def init_worker_state(self, worker_weights: gradient_learner.WorkerWeights,
                         key: PRNGKey) -> UnrollState:
     return UnrollState()
 
-  def compute_gradient_estimate(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def compute_gradient_estimate(  # pyrefly: ignore[bad-override]
       self,
       worker_weights,
       key,

learned_optimization/outer_trainers/gradient_learner.py

@@ -285,7 +285,7 @@ def update(
       metrics["none||theta_grad_norm"] = tree_utils.tree_norm(grads)
       metrics["none||theta_grad_abs_mean"] = tree_utils.tree_mean_abs(grads)
 
-    return GradientLearnerState(theta_opt_state), metrics  # pytype: disable=bad-return-type  # jax-ndarray
+    return GradientLearnerState(theta_opt_state), metrics
 
 
 class GradientEstimator(abc.ABC):