Add tree_bytes function.

Author: carlosgmartin

docs/api/utilities.rst

@@ -119,6 +119,7 @@ Tree
     tree_scale
     tree_set
     tree_size
+    tree_bytes
     tree_sub
     tree_sum
     tree_vdot
@@ -221,6 +222,10 @@ Tree size
 ~~~~~~~~~
 .. autofunction:: tree_size
 
+Tree bytes
+~~~~~~~~~
+.. autofunction:: tree_bytes
+
 Tree subtract
 ~~~~~~~~~~~~~
 .. autofunction:: tree_sub

optax/contrib/_sophia.py

@@ -161,9 +161,8 @@ def update_fn(updates, state: SophiaState, params=None, **hess_fn_kwargs):
           lambda x, y: x + y,
           jax.tree.map(lambda u: jnp.sum(jnp.abs(u) < clip_threshold), updates),
       )
-      total_tree_size = sum(x.size for x in jax.tree.leaves(updates))
       if verbose:
-        win_rate = sum_not_clipped / total_tree_size
+        win_rate = sum_not_clipped / optax.tree.size(updates)
         jax.lax.cond(
             count_inc % print_win_rate_every_n_steps == 0,
             lambda: jax.debug.print("Sophia optimizer win rate: {}", win_rate),

optax/tree/__init__.py

@@ -47,6 +47,7 @@
 real = _tree_math.tree_real
 scale = _tree_math.tree_scale
 size = _tree_math.tree_size
+bytes = _tree_math.tree_bytes  # pylint: disable=redefined-builtin
 sub = _tree_math.tree_sub
 sum = _tree_math.tree_sum  # pylint: disable=redefined-builtin
 update_infinity_moment = _tree_math.tree_update_infinity_moment

optax/tree_utils/__init__.py

@@ -46,6 +46,7 @@
 from optax.tree_utils._tree_math import tree_real
 from optax.tree_utils._tree_math import tree_scale
 from optax.tree_utils._tree_math import tree_size
+from optax.tree_utils._tree_math import tree_bytes
 from optax.tree_utils._tree_math import tree_sub
 from optax.tree_utils._tree_math import tree_sum
 from optax.tree_utils._tree_math import tree_update_infinity_moment

optax/tree_utils/_tree_math.py

@@ -209,6 +209,18 @@ def tree_size(tree: Any) -> int:
   return sum(jnp.size(leaf) for leaf in jax.tree.leaves(tree))
 
 
+def tree_bytes(tree: Any) -> int:
+  r"""Total number of bytes in a pytree.
+
+  Args:
+    tree: pytree
+
+  Returns:
+    the total number of bytes in the pytree.
+  """
+  return sum(jnp.asarray(leaf).nbytes for leaf in jax.tree.leaves(tree))
+
+
 def tree_conj(tree: Any) -> Any:
   """Compute the conjugate of a pytree.