Fix global_norm() signature

google-deepmind · Feb 7, 2023 · 3d0c422 · 3d0c422
1 parent a77d69c
commit 3d0c422
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)))