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,23 @@ 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 size of the pytree in bytes.
+
+  .. warning::
+    It is assumed that every leaf's dtype has an integer byte size.
+    Fractional byte sizes may yield an incorrect result.
+    For example, ``int4`` might be only half a byte on device.
+  """
+  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.
 

optax/tree_utils/_tree_math_test.py

@@ -172,6 +172,13 @@ def test_tree_size(self, key):
     got = tu.tree_size(tree)
     np.testing.assert_allclose(expected, got)
 
+  @parameterized.product(
+    n=[1, 10, 100, 1000],
+    dtype=[jnp.int16, jnp.int32, jnp.float16, jnp.float32],
+  )
+  def test_tree_bytes(self, n, dtype):
+    assert tu.tree_bytes(jnp.ones(n, dtype)) == n * dtype.dtype.itemsize
+
   def test_tree_conj(self):
     expected = jnp.conj(self.array_a)
     got = tu.tree_conj(self.array_a)