No public description

PiperOrigin-RevId: 574170487
google · Oct 17, 2023 · 1f9823d · 1f9823d
1 parent 463ab9a
commit 1f9823d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/learned_optimization/outer_train.py b/learned_optimization/outer_train.py
@@ -235,7 +235,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
+  return metrics, worker_ids, applied_inner_steps  # pytype: disable=bad-return-type
 
 
 def maybe_resample_gradient_estimators(

diff --git 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
+  return jnp.mean((pos_loss + neg_loss) / 2.0), es_grad  # pytype: disable=bad-return-type
 
 
 @gin.configurable