Merge pull request #481 from brentyi:fix_global_norm_sig

PiperOrigin-RevId: 507720585
google-deepmind · Feb 7, 2023 · 4919451 · 4919451
2 parents a77d69c + 3d0c422
commit 4919451
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optax/_src/linear_algebra.py b/optax/_src/linear_algebra.py
@@ -24,7 +24,7 @@
 from optax._src import numerics
 
 
-def global_norm(updates: base.Updates) -> base.Updates:
+def global_norm(updates: base.PyTree) -> chex.Array:
   """Compute the global norm across a nested structure of tensors."""
   return jnp.sqrt(sum(
       jnp.sum(numerics.abs_sq(x)) for x in jax.tree_util.tree_leaves(updates)))