Merge pull request #1321 from carlosgmartin:tree_utils

PiperOrigin-RevId: 773082934

Author: OptaxDev

docs/api/utilities.rst

@@ -110,13 +110,15 @@ Tree
     tree_norm
     tree_map_params
     tree_max
+    tree_min
     tree_mul
     tree_ones_like
     tree_random_like
     tree_real
     tree_split_key_like
     tree_scale
     tree_set
+    tree_size
     tree_sub
     tree_sum
     tree_vdot
@@ -183,6 +185,10 @@ Tree max
 ~~~~~~~~
 .. autofunction:: tree_max
 
+Tree min
+~~~~~~~~
+.. autofunction:: tree_min
+
 Tree multiply
 ~~~~~~~~~~~~~
 .. autofunction:: tree_mul
@@ -211,6 +217,10 @@ Set values in a tree
 ~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: tree_set
 
+Tree size
+~~~~~~~~~
+.. autofunction:: tree_size
+
 Tree subtract
 ~~~~~~~~~~~~~
 .. autofunction:: tree_sub

optax/tree/__init__.py

@@ -40,11 +40,13 @@
 div = _tree_math.tree_div
 full_like = _tree_math.tree_full_like
 max = _tree_math.tree_max  # pylint: disable=redefined-builtin
+min = _tree_math.tree_min  # pylint: disable=redefined-builtin
 mul = _tree_math.tree_mul
 norm = _tree_math.tree_norm
 ones_like = _tree_math.tree_ones_like
 real = _tree_math.tree_real
 scale = _tree_math.tree_scale
+size = _tree_math.tree_size
 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

@@ -39,11 +39,13 @@
 from optax.tree_utils._tree_math import tree_div
 from optax.tree_utils._tree_math import tree_full_like
 from optax.tree_utils._tree_math import tree_max
+from optax.tree_utils._tree_math import tree_min
 from optax.tree_utils._tree_math import tree_mul
 from optax.tree_utils._tree_math import tree_norm
 from optax.tree_utils._tree_math import tree_ones_like
 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_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

@@ -183,6 +183,32 @@ def tree_max(tree: Any) -> chex.Numeric:
   return jax.tree.reduce(jnp.maximum, maxes, initializer=jnp.array(-jnp.inf))
 
 
+def tree_min(tree: Any) -> chex.Numeric:
+  """Compute the min of all the elements in a pytree.
+
+  Args:
+    tree: pytree.
+
+  Returns:
+    a scalar value.
+  """
+  mins = jax.tree.map(jnp.min, tree)
+  # initializer=jnp.inf should work but pytype wants a jax.Array.
+  return jax.tree.reduce(jnp.minimum, mins, initializer=jnp.array(jnp.inf))
+
+
+def tree_size(tree: Any) -> int:
+  r"""Total size of a pytree.
+
+  Args:
+    tree: pytree
+
+  Returns:
+    the total size of the pytree.
+  """
+  return sum(jnp.size(leaf) 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

@@ -152,6 +152,26 @@ def test_tree_max(self, key):
     got = tu.tree_max(tree)
     np.testing.assert_allclose(expected, got)
 
+  @parameterized.parameters(
+      'array_a', 'tree_a', 'tree_a_dict', 'tree_b', 'tree_b_dict'
+  )
+  def test_tree_min(self, key):
+    tree = self.data[key]
+    values, _ = flatten_util.ravel_pytree(tree)
+    expected = jnp.min(values)
+    got = tu.tree_min(tree)
+    np.testing.assert_allclose(expected, got)
+
+  @parameterized.parameters(
+      'array_a', 'tree_a', 'tree_a_dict', 'tree_b', 'tree_b_dict'
+  )
+  def test_tree_size(self, key):
+    tree = self.data[key]
+    values, _ = flatten_util.ravel_pytree(tree)
+    expected = jnp.size(values)
+    got = tu.tree_size(tree)
+    np.testing.assert_allclose(expected, got)
+
   def test_tree_conj(self):
     expected = jnp.conj(self.array_a)
     got = tu.tree_conj(self.array_a)