From 64f83138777cb923e4a83f6cf485b075037bf2bd Mon Sep 17 00:00:00 2001 From: Oleh Prypin Date: Wed, 7 Oct 2026 06:33:22 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 995070179 --- learned_optimization/circular_buffer.py | 9 +++-- learned_optimization/distributed.py | 6 ++- learned_optimization/eval_training.py | 4 +- learned_optimization/filesystem.py | 2 +- learned_optimization/jax_utils.py | 2 +- .../learned_optimizers/adafac_mlp_lopt.py | 8 ++-- .../learned_optimizers/adafac_nominal.py | 8 ++-- .../learned_optimizers/base.py | 4 +- .../learned_optimizers/common.py | 4 +- .../learned_optimizers/mlp_lopt.py | 2 +- .../learned_optimizers/nn_adam.py | 29 ++++++-------- .../learned_optimizers/opt_from_checkpoint.py | 6 +-- .../learned_optimizers/rnn_mlp_lopt.py | 2 +- learned_optimization/outer_train.py | 8 ++-- learned_optimization/outer_trainers/common.py | 11 ++++-- .../outer_trainers/full_es.py | 2 +- .../outer_trainers/full_grad.py | 2 +- .../outer_trainers/gradient_learner.py | 2 +- .../outer_trainers/lopt_truncated_step.py | 39 ++++++++++--------- .../outer_trainers/test_utils.py | 2 +- .../outer_trainers/truncated_es.py | 2 +- .../outer_trainers/truncated_pes.py | 4 +- .../outer_trainers/truncation_schedule.py | 2 +- learned_optimization/summary.py | 12 +++--- .../tasks/parametric/cfgobject.py | 14 +++---- .../tasks/parametric/image_mlp.py | 2 +- .../tasks/parametric/image_resnet.py | 11 ++++-- .../tasks/parametric/parametric_utils.py | 4 +- learned_optimization/tree_utils.py | 7 ++-- 29 files changed, 110 insertions(+), 100 deletions(-) diff --git a/learned_optimization/circular_buffer.py b/learned_optimization/circular_buffer.py index 371f6864..a9c543a3 100644 --- a/learned_optimization/circular_buffer.py +++ b/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 diff --git a/learned_optimization/distributed.py b/learned_optimization/distributed.py index a99d0db3..0f4a2bf3 100644 --- a/learned_optimization/distributed.py +++ b/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.""" diff --git a/learned_optimization/eval_training.py b/learned_optimization/eval_training.py index 74a8d9ab..5840ba9b 100644 --- a/learned_optimization/eval_training.py +++ b/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 diff --git a/learned_optimization/filesystem.py b/learned_optimization/filesystem.py index be3b96df..8e8663e8 100644 --- a/learned_optimization/filesystem.py +++ b/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] diff --git a/learned_optimization/jax_utils.py b/learned_optimization/jax_utils.py index d6077b21..7f7376d1 100644 --- a/learned_optimization/jax_utils.py +++ b/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): diff --git a/learned_optimization/learned_optimizers/adafac_mlp_lopt.py b/learned_optimization/learned_optimizers/adafac_mlp_lopt.py index 8733a001..c327f41f 100644 --- a/learned_optimization/learned_optimizers/adafac_mlp_lopt.py +++ b/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, diff --git a/learned_optimization/learned_optimizers/adafac_nominal.py b/learned_optimization/learned_optimizers/adafac_nominal.py index da304996..99d95602 100644 --- a/learned_optimization/learned_optimizers/adafac_nominal.py +++ b/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, diff --git a/learned_optimization/learned_optimizers/base.py b/learned_optimization/learned_optimizers/base.py index 8df331e5..f1c045c3 100644 --- a/learned_optimization/learned_optimizers/base.py +++ b/learned_optimization/learned_optimizers/base.py @@ -214,7 +214,7 @@ 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]) @@ -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): # 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) diff --git a/learned_optimization/learned_optimizers/common.py b/learned_optimization/learned_optimizers/common.py index 762273df..bcebcac8 100644 --- a/learned_optimization/learned_optimizers/common.py +++ b/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) diff --git a/learned_optimization/learned_optimizers/mlp_lopt.py b/learned_optimization/learned_optimizers/mlp_lopt.py index 72790e97..a75ac893 100644 --- a/learned_optimization/learned_optimizers/mlp_lopt.py +++ b/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, diff --git a/learned_optimization/learned_optimizers/nn_adam.py b/learned_optimization/learned_optimizers/nn_adam.py index b3a01eca..3eb3260b 100644 --- a/learned_optimization/learned_optimizers/nn_adam.py +++ b/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, diff --git a/learned_optimization/learned_optimizers/opt_from_checkpoint.py b/learned_optimization/learned_optimizers/opt_from_checkpoint.py index bba60931..588c32cf 100644 --- a/learned_optimization/learned_optimizers/opt_from_checkpoint.py +++ b/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) diff --git a/learned_optimization/learned_optimizers/rnn_mlp_lopt.py b/learned_optimization/learned_optimizers/rnn_mlp_lopt.py index 25bf7b23..8c1ef27f 100644 --- a/learned_optimization/learned_optimizers/rnn_mlp_lopt.py +++ b/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] ) diff --git a/learned_optimization/outer_train.py b/learned_optimization/outer_train.py index aef71027..7c9ed960 100644 --- a/learned_optimization/outer_train.py +++ b/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, diff --git a/learned_optimization/outer_trainers/common.py b/learned_optimization/outer_trainers/common.py index c0d1ec16..2df4ac56 100644 --- a/learned_optimization/outer_trainers/common.py +++ b/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, diff --git a/learned_optimization/outer_trainers/full_es.py b/learned_optimization/outer_trainers/full_es.py index aea4104c..87a355ed 100644 --- a/learned_optimization/outer_trainers/full_es.py +++ b/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 diff --git a/learned_optimization/outer_trainers/full_grad.py b/learned_optimization/outer_trainers/full_grad.py index d26765ca..4f377547 100644 --- a/learned_optimization/outer_trainers/full_grad.py +++ b/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, diff --git a/learned_optimization/outer_trainers/gradient_learner.py b/learned_optimization/outer_trainers/gradient_learner.py index a61b13f1..c23f16bd 100644 --- a/learned_optimization/outer_trainers/gradient_learner.py +++ b/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): diff --git a/learned_optimization/outer_trainers/lopt_truncated_step.py b/learned_optimization/outer_trainers/lopt_truncated_step.py index 47f66b6b..44cdae0e 100644 --- a/learned_optimization/outer_trainers/lopt_truncated_step.py +++ b/learned_optimization/outer_trainers/lopt_truncated_step.py @@ -83,7 +83,7 @@ def train(unroll_state): params = opt.get_params(unroll_state) loss, grad = jax.value_and_grad(self.task.loss)(params, key, data) unroll_state = opt.update(unroll_state, grad, loss=loss) - out = truncated_step.TruncatedUnrollOut( # pytype: disable=wrong-arg-types # jax-ndarray + out = truncated_step.TruncatedUnrollOut( loss=loss, is_done=False, # pyrefly: ignore[bad-argument-type] task_param=None, @@ -95,7 +95,7 @@ def train(unroll_state): def reset(unroll_state): params = self.task.init(key) unroll_state = self.lopt.opt_fn(theta).init(params) - out = truncated_step.TruncatedUnrollOut( # pytype: disable=wrong-arg-types # jax-ndarray + out = truncated_step.TruncatedUnrollOut( loss=0.0, # pyrefly: ignore[bad-argument-type] is_done=True, # pyrefly: ignore[bad-argument-type] task_param=None, @@ -111,7 +111,7 @@ def meta_loss_batch(self, theta: MetaParams, unroll_state: InnerState, key: chex.PRNGKey, data: OuterBatch, outer_state: OuterState) -> jnp.ndarray: params = self.lopt.opt_fn(theta).get_params(unroll_state) - return self.task.loss(params, key, data) # pytype: disable=bad-return-type # jax-ndarray + return self.task.loss(params, key, data) # pyrefly: ignore[bad-return] def get_batch(self, steps: Optional[int] = None) -> InnerBatch: return training.vec_get_batch(self.task, steps, split="train") @@ -196,7 +196,7 @@ def _init_truncation_state( theta, is_training=True).init( inner_param, inner_state, num_steps=num_steps, key=key4) - return TruncatedUnrollState( # pytype: disable=wrong-arg-types # jax-ndarray + return TruncatedUnrollState( inner_opt_state=opt_state, inner_step=jnp.asarray(0, dtype=jnp.int32), truncation_state=trunc_state, @@ -220,7 +220,7 @@ def progress_or_reset_inner_opt_state( meta_loss_with_aux_key: Optional[str] = None, ) -> Tuple[T, G, int, jnp.ndarray]: """Train a single step, or reset the current inner problem.""" - # summary.summary("num_steps", num_steps, aggregation="sample") # pytype: disable=wrong-arg-types # jax-ndarray + # summary.summary("num_steps", num_steps, aggregation="sample") def true_fn(key): """Reset the state of the inner-problem.""" @@ -234,7 +234,7 @@ def true_fn(key): p, s = task_family.task_fn(task_param).init_with_state(key2) next_inner_opt_state = opt.init(p, s, num_steps=num_steps, key=key3) - # summary.summary("opt_init_num_steps", num_steps) # pytype: disable=wrong-arg-types # jax-ndarray + # summary.summary("opt_init_num_steps", num_steps) return next_inner_opt_state, task_param, jnp.asarray(0), jnp.asarray(0.) @@ -280,8 +280,9 @@ def loss_fn(p, s, key, data): return next_inner_opt_state, task_param, next_inner_step, jnp.asarray( meta_loss, dtype=jnp.float32) - next_inner_opt_state, task_param, next_inner_step, meta_loss = cond_fn( # pytype: disable=wrong-arg-types # jax-types - jnp.logical_not(is_done), false_fn, true_fn, key) + next_inner_opt_state, task_param, next_inner_step, meta_loss = cond_fn( + jnp.logical_not(is_done), false_fn, true_fn, key # pyrefly: ignore[bad-argument-type] + ) return next_inner_opt_state, task_param, next_inner_step, meta_loss @@ -298,7 +299,7 @@ def vectorized_loss_and_aux(task_family: tasks_base.TaskFamily, opt = learned_opt.opt_fn(theta, is_training=True) p, s = opt.get_params_state(inner_opt_state) l, _, aux = task.loss_with_state_and_aux(p, s, key, data) - return l, aux # pytype: disable=bad-return-type # jax-ndarray + return l, aux # pyrefly: ignore[bad-return] def _truncated_unroll_one_step( @@ -322,15 +323,15 @@ def _truncated_unroll_one_step( num_steps = state.truncation_state.length next_inner_opt_state, task_param, next_inner_step, l = ( - progress_or_reset_inner_opt_state( # pytype: disable=wrong-arg-types # jax-ndarray + progress_or_reset_inner_opt_state( task_family=task_family, opt=learned_opt.opt_fn(theta), num_steps=num_steps, key=key1, inner_opt_state=state.inner_opt_state, task_param=state.task_param, - inner_step=state.inner_step, - is_done=state.is_done, + inner_step=state.inner_step, # pyrefly: ignore[bad-argument-type] + is_done=state.is_done, # pyrefly: ignore[bad-argument-type] data=data, meta_loss_with_aux_key=meta_loss_with_aux_key, ) @@ -343,19 +344,19 @@ def _truncated_unroll_one_step( # opt = learned_opt.opt_fn(theta, is_training=sTrue) # summary.summarize_inner_params(opt.get_params(next_inner_opt_state)) - output_state = TruncatedUnrollState( # pytype: disable=wrong-arg-types # jax-ndarray + output_state = TruncatedUnrollState( inner_opt_state=next_inner_opt_state, - inner_step=next_inner_step, + inner_step=next_inner_step, # pyrefly: ignore[bad-argument-type] truncation_state=next_truncation_state, task_param=task_param, - is_done=is_done, + is_done=is_done, # pyrefly: ignore[bad-argument-type] ) - out = truncated_step.TruncatedUnrollOut( # pytype: disable=wrong-arg-types # jax-ndarray - is_done=is_done, + out = truncated_step.TruncatedUnrollOut( + is_done=is_done, # pyrefly: ignore[bad-argument-type] loss=l, - mask=(next_inner_step != 0), - iteration=next_inner_step, + mask=(next_inner_step != 0), # pyrefly: ignore[bad-argument-type] + iteration=next_inner_step, # pyrefly: ignore[bad-argument-type] task_param=state.task_param, ) diff --git a/learned_optimization/outer_trainers/test_utils.py b/learned_optimization/outer_trainers/test_utils.py index fc2caab0..1b2644ff 100644 --- a/learned_optimization/outer_trainers/test_utils.py +++ b/learned_optimization/outer_trainers/test_utils.py @@ -22,7 +22,7 @@ def trainer_smoketest(trainer): key = jax.random.PRNGKey(0) theta = trainer.truncated_step.outer_init(key) worker_weights = gradient_learner.WorkerWeights( - theta, None, gradient_learner.OuterState(1) # pytype: disable=wrong-arg-types # jax-ndarray + theta, None, gradient_learner.OuterState(1) # pyrefly: ignore[bad-argument-type] ) state = trainer.init_worker_state(worker_weights, key=key) out, metrics = trainer.compute_gradient_estimate( diff --git a/learned_optimization/outer_trainers/truncated_es.py b/learned_optimization/outer_trainers/truncated_es.py index a0fa3d76..3c9ca129 100644 --- a/learned_optimization/outer_trainers/truncated_es.py +++ b/learned_optimization/outer_trainers/truncated_es.py @@ -85,7 +85,7 @@ def flat_first(x): 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, p_ys, delta_loss # pytype: disable=bad-return-type + return jnp.mean((pos_loss + neg_loss) / 2.0), es_grad, p_ys, delta_loss # pyrefly: ignore[bad-return] @gin.configurable diff --git a/learned_optimization/outer_trainers/truncated_pes.py b/learned_optimization/outer_trainers/truncated_pes.py index fe5b2c80..1b20bbca 100644 --- a/learned_optimization/outer_trainers/truncated_pes.py +++ b/learned_optimization/outer_trainers/truncated_pes.py @@ -136,7 +136,7 @@ def _switch_one_accum(a, b): new_accumulator, p_ys, delta_losses, - ) # pytype: disable=bad-return-type + ) @gin.configurable @@ -202,7 +202,7 @@ def get_datas(self): ] @profile.wrap() - def compute_gradient_estimate( # pytype: disable=signature-mismatch # overriding-parameter-type-checks + def compute_gradient_estimate( self, worker_weights: gradient_learner.WorkerWeights, key: PRNGKey, diff --git a/learned_optimization/outer_trainers/truncation_schedule.py b/learned_optimization/outer_trainers/truncation_schedule.py index 62e9682c..3f011659 100644 --- a/learned_optimization/outer_trainers/truncation_schedule.py +++ b/learned_optimization/outer_trainers/truncation_schedule.py @@ -53,7 +53,7 @@ def __init__(self, total_length: int): self._total_length = total_length def init(self, key: PRNGKey, outer_state: Any) -> ConstantTruncationState: - return ConstantTruncationState(length=self._total_length) # pytype: disable=wrong-arg-types # jax-ndarray + return ConstantTruncationState(length=self._total_length) # pyrefly: ignore[bad-argument-type] def next_state( self, state: ConstantTruncationState, step: int, key: PRNGKey, diff --git a/learned_optimization/summary.py b/learned_optimization/summary.py index 355e406c..62c0cc19 100644 --- a/learned_optimization/summary.py +++ b/learned_optimization/summary.py @@ -178,7 +178,7 @@ def aggregate_metric_list( keys = [None] * len(all_metrics) for ki, (k, vs) in enumerate(all_metrics.items()): metrics[k] = aggregate_metric(k, vs, use_jnp=use_jnp, key=keys[ki]) - return metrics # pytype: disable=bad-return-type # jax-ndarray + return metrics def aggregate_metric(k: str, @@ -195,7 +195,7 @@ def aggregate_metric(k: str, if agg == AggregationType.mean: # size is known at compile time. size = onp.sum([onp.prod(v.shape) for v in vs]) - return xnp.sum(xnp.asarray([xnp.sum(v) / size for v in vs])) # pytype: disable=bad-return-type # jnp-type + return xnp.sum(xnp.asarray([xnp.sum(v) / size for v in vs])) # pyrefly: ignore[bad-return] elif agg == AggregationType.sample: vs = xnp.concatenate([xnp.asarray(v).ravel() for v in vs], axis=0) # pyrefly: ignore[bad-assignment] if use_jnp: @@ -204,13 +204,13 @@ def aggregate_metric(k: str, else: i = onp.random.randint(0, len(vs), dtype=xnp.int32) - return vs[i] # pytype: disable=bad-return-type # jnp-type + return vs[i] # pyrefly: ignore[bad-index] elif agg == AggregationType.collect: # This might be multi dim if vmap is used, so ravel first. - return xnp.concatenate([xnp.asarray(v).ravel() for v in vs], axis=0) # pytype: disable=bad-return-type # jnp-type + return xnp.concatenate([xnp.asarray(v).ravel() for v in vs], axis=0) # pyrefly: ignore[bad-return] elif agg == AggregationType.tensor: assert len(vs) == 1 - return vs[0] # pytype: disable=bad-return-type # jnp-type + return vs[0] # pyrefly: ignore[bad-return] elif agg == AggregationType.none: if len(vs) != 1: raise ValueError("when using no aggregation one must ensure only scalar " @@ -220,7 +220,7 @@ def aggregate_metric(k: str, if val.size != 1: raise ValueError("Value with none aggregation type was not a scalar?" f" Found {val}") - return xnp.reshape(val, ()) # pytype: disable=bad-return-type # jnp-type + return xnp.reshape(val, ()) # pyrefly: ignore[bad-return] else: raise ValueError(f"Unsupported Aggregation type {agg}") diff --git a/learned_optimization/tasks/parametric/cfgobject.py b/learned_optimization/tasks/parametric/cfgobject.py index b4769ff7..ef9b472a 100644 --- a/learned_optimization/tasks/parametric/cfgobject.py +++ b/learned_optimization/tasks/parametric/cfgobject.py @@ -216,20 +216,20 @@ def featurize_value(val: Any) -> jnp.ndarray: # preprocess some of the tags heuristically. empty = jnp.zeros((8,)) if isinstance(val, str): - return empty, jnp.asarray(hash_trick(val)) # pytype: disable=bad-return-type # jax-ndarray + return empty, jnp.asarray(hash_trick(val)) # pyrefly: ignore[bad-return] if isinstance(val, Sequence): - return pad(val), 0 # pytype: disable=bad-return-type # jax-ndarray + return pad(val), 0 # pyrefly: ignore[bad-argument-type, bad-return] if isinstance(val, (int, onp.int32, onp.int64)): - return empty, jnp.asarray(val) # pytype: disable=bad-return-type # jax-ndarray + return empty, jnp.asarray(val) # pyrefly: ignore[bad-return] if isinstance(val, (bool, onp.bool_)): return pad(onp.asarray(val, dtype=onp.float32)), jnp.asarray( # type: ignore # jax-ndarray val, dtype=onp.int32) if isinstance(val, (float, onp.float32)): - return pad(val), 0 # pytype: disable=bad-return-type # jax-ndarray + return pad(val), 0 # pyrefly: ignore[bad-argument-type, bad-return] if isinstance(val, jnp.ndarray): - return pad(val), 0 # pytype: disable=bad-return-type # jax-ndarray + return pad(val), 0 # pyrefly: ignore[bad-return] if isinstance(val, onp.ndarray): - return pad(val), 0 # pytype: disable=bad-return-type # jax-ndarray + return pad(val), 0 # pyrefly: ignore[bad-argument-type, bad-return] raise NotImplementedError(val, type(val)) @@ -241,7 +241,7 @@ def featurize_cfg_path(path: str) -> jnp.ndarray: continue idxs.append(hash_trick(s)) unpadded = onp.asarray(idxs) - return pad(unpadded) # pytype: disable=wrong-arg-types # jax-ndarray + return pad(unpadded) # pyrefly: ignore[bad-argument-type] def featurize( diff --git a/learned_optimization/tasks/parametric/image_mlp.py b/learned_optimization/tasks/parametric/image_mlp.py index 8824b407..517a0471 100644 --- a/learned_optimization/tasks/parametric/image_mlp.py +++ b/learned_optimization/tasks/parametric/image_mlp.py @@ -86,7 +86,7 @@ def init(self, rng: PRNGKey) -> Params: image = next(self.datasets.train)["image"] # pyrefly: ignore[missing-attribute] return init_net(rng, image) - def loss(self, params: Params, rng: PRNGKey, data: Batch) -> jnp.ndarray: # pytype: disable=signature-mismatch # jax-ndarray + def loss(self, params: Params, rng: PRNGKey, data: Batch) -> jnp.ndarray: unused_init_net, apply_net = hk.without_apply_rng( hk.transform(_forward)) diff --git a/learned_optimization/tasks/parametric/image_resnet.py b/learned_optimization/tasks/parametric/image_resnet.py index 4ea9e347..7c0a7a27 100644 --- a/learned_optimization/tasks/parametric/image_resnet.py +++ b/learned_optimization/tasks/parametric/image_resnet.py @@ -90,15 +90,20 @@ def init_with_state(self, key: PRNGKey) -> Tuple[Params, ModelState]: params, state = init_net(key, image) return params, state - def loss_with_state(self, params: Params, state: ModelState, key: PRNGKey, # pytype: disable=signature-mismatch # jax-ndarray - data: Batch) -> jnp.ndarray: + def loss_with_state( # pyrefly: ignore[bad-override] + self, + params: Params, + state: ModelState, + key: PRNGKey, + data: Batch, + ) -> jnp.ndarray: net = hk.transform_with_state(_forward) image = data["image"] logits, state = net.apply(params, state, key, image) labels = jax.nn.one_hot(data["label"], num_classes) vec_loss = base.softmax_cross_entropy(logits=logits, labels=labels) - return jnp.mean(vec_loss), state # pytype: disable=bad-return-type # jax-ndarray + return jnp.mean(vec_loss), state # pyrefly: ignore[bad-return] def loss_with_state_and_aux( self, params: Params, state: ModelState, key: PRNGKey, data: Batch diff --git a/learned_optimization/tasks/parametric/parametric_utils.py b/learned_optimization/tasks/parametric/parametric_utils.py index 3289697a..6dcb8590 100644 --- a/learned_optimization/tasks/parametric/parametric_utils.py +++ b/learned_optimization/tasks/parametric/parametric_utils.py @@ -32,7 +32,7 @@ def choice(key, vals): - return vals[jax.random.randint(key, [], 0, len(vals))] # pytype: disable=unsupported-operands # jax-types + return vals[jax.random.randint(key, [], 0, len(vals))] def log_int(key, lower, upper): @@ -94,7 +94,7 @@ class SampleActivation: @classmethod def sample(cls, key) -> Mapping[str, int]: - return {"index": jax.random.randint(key, [], 0, len(cls.acts))} # pytype: disable=bad-return-type # jax-types + return {"index": jax.random.randint(key, [], 0, len(cls.acts))} # pyrefly: ignore[bad-assignment] @classmethod def get_dynamic(cls, act): diff --git a/learned_optimization/tree_utils.py b/learned_optimization/tree_utils.py index d346f029..a675ad43 100644 --- a/learned_optimization/tree_utils.py +++ b/learned_optimization/tree_utils.py @@ -142,8 +142,9 @@ def map_named(function: Callable[[str, Any], Any], Struct with the same pytree. """ if isinstance(val, Mapping): - return type(val)( # pyrefly: ignore[bad-instantiation] - **{k: map_named(function, v, key + "/" + k) for k, v in val.items()}) # pyrefly: ignore[unsupported-operation] + return type(val)( + **{k: map_named(function, v, key + "/" + k) for k, v in val.items()} # pyrefly: ignore[unsupported-operation] + ) elif isinstance(val, tuple) or isinstance(val, list): return type(val)( * @@ -205,7 +206,7 @@ def get_name(k, v): del v return k - keys = jax.tree_util.tree_leaves(map_named(get_name, "", values)) # pytype: disable=wrong-arg-types # numpy-scalars + keys = jax.tree_util.tree_leaves(map_named(get_name, "", values)) # pyrefly: ignore[bad-argument-type] keys = [str(i) for i, v in enumerate(vals)] if not strict: functions = list(functions) + [lambda k, v: True]