Internal change

PiperOrigin-RevId: 530350759

Author: rchen152

learned_optimization/learned_optimizers/adafac_mlp_lopt.py

@@ -404,14 +404,15 @@ def init(
             iteration=jnp.asarray(0, dtype=jnp.int32),
             num_steps=jnp.asarray(num_steps))
 
-      def update(self,
-                 opt_state: AdafacMLPLOptState,
-                 grad: opt_base.Gradient,
-                 loss: jnp.ndarray,
-                 model_state: Optional[opt_base.ModelState] = None,
-                 is_valid: bool = False,
-                 key: Optional[PRNGKey] = None) -> AdafacMLPLOptState:
-
+      def update(
+          self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+          opt_state: AdafacMLPLOptState,
+          grad: opt_base.Gradient,
+          loss: jnp.ndarray,
+          model_state: Optional[opt_base.ModelState] = None,
+          is_valid: bool = False,
+          key: Optional[PRNGKey] = None,
+      ) -> AdafacMLPLOptState:
         mom_roll, rms_roll, fac_vec_roll = self._get_rolling()
         next_mom_rolling = mom_roll.update(opt_state.mom_rolling, grad)
         next_rms_rolling = rms_roll.update(opt_state.rms_rolling, grad)

learned_optimization/learned_optimizers/adafac_nominal.py

@@ -477,14 +477,15 @@ def init(
             iteration=jnp.asarray(0, dtype=jnp.int32),
             num_steps=jnp.asarray(num_steps))
 
-      def update(self,
-                 opt_state: AdafacMLPLOptState,
-                 grad: opt_base.Gradient,
-                 loss: jnp.ndarray,
-                 model_state: Optional[opt_base.ModelState] = None,
-                 is_valid: bool = False,
-                 key: Optional[PRNGKey] = None) -> AdafacMLPLOptState:
-
+      def update(
+          self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+          opt_state: AdafacMLPLOptState,
+          grad: opt_base.Gradient,
+          loss: jnp.ndarray,
+          model_state: Optional[opt_base.ModelState] = None,
+          is_valid: bool = False,
+          key: Optional[PRNGKey] = None,
+      ) -> AdafacMLPLOptState:
         mom_roll, rms_roll, fac_vec_roll = self._get_rolling()
         next_mom_rolling = mom_roll.update(opt_state.mom_rolling, grad)
         next_rms_rolling = rms_roll.update(opt_state.rms_rolling, grad)

learned_optimization/learned_optimizers/base.py

@@ -222,7 +222,7 @@ def get_params(self, state):
   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):
+  def update(self, opt_state, grad, model_state=None, **kwargs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     # apply to both opts
     new_opt_states = [
         opt.update(os, grad, model_state=model_state, **kwargs)

learned_optimization/learned_optimizers/mlp_lopt.py

@@ -115,14 +115,15 @@ def init(self,
             rolling_features=common.vec_rolling_mom(decays).init(params),
             iteration=jnp.asarray(0, dtype=jnp.int32))
 
-      def update(self,
-                 opt_state: MLPLOptState,
-                 grad: Any,
-                 loss: float,
-                 model_state: Any = None,
-                 is_valid: bool = False,
-                 key: Optional[PRNGKey] = None) -> MLPLOptState:
-
+      def update(
+          self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+          opt_state: MLPLOptState,
+          grad: Any,
+          loss: float,
+          model_state: Any = None,
+          is_valid: bool = False,
+          key: Optional[PRNGKey] = None,
+      ) -> MLPLOptState:
         next_rolling_features = common.vec_rolling_mom(decays).update(
             opt_state.rolling_features, grad)
 

learned_optimization/optimizers/base.py

@@ -140,7 +140,7 @@ def init(self, params, model_state=None, num_steps=None, **kwargs):
         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):
+  def update(self, opt_state, grad, model_state=None, **kwargs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     base_params = opt_state.params
 
     next_mag_opt_state = self.magnitude_opt.update(

learned_optimization/optimizers/learning_rate_schedules.py

@@ -77,8 +77,7 @@ def __init__(self,
     self.constant_fraction = constant_fraction
     self.warmup_fraction = warmup_fraction
 
-  def __call__(self, global_step, max_steps) -> chex.Array:
-
+  def __call__(self, global_step, max_steps) -> chex.Array:  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     def fload32(x):
       """Convert input to float32."""
       return jnp.asarray(x, dtype=onp.float32)

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(
+  def compute_gradient_estimate(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
       self,
       worker_weights,
       key,

learned_optimization/research/hysteresis/truncated_es_shared_noise.py

@@ -97,15 +97,13 @@ def init_worker_state(self, worker_weights: gradient_learner.WorkerWeights,
         epsilons=epsilons)
 
   @profile.wrap()
-  def compute_gradient_estimate(
+  def compute_gradient_estimate(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
       self,
       worker_weights,
       key: PRNGKey,
-      state:
-      TruncatedESSharedNoiseAttributes,  # this is the same state returned by init_worker_state
+      state: TruncatedESSharedNoiseAttributes,  # this is the same state returned by init_worker_state
       with_summary=False,
   ) -> Tuple[gradient_learner.GradientEstimatorOut, Mapping[str, jnp.ndarray]]:
-
     # because we have a for loop we let haiku manages the key
     rng = hk.PRNGSequence(key)