Add tree utility functions.

Author: carlosgmartin

docs/api/utilities.rst

@@ -110,6 +110,8 @@ Tree
     tree_norm
     tree_map_params
     tree_max
+    tree_min
+    tree_size
     tree_mul
     tree_ones_like
     tree_random_like

optax/tree/__init__.py

@@ -40,6 +40,8 @@
 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
+size = _tree_math.tree_size
 mul = _tree_math.tree_mul
 norm = _tree_math.tree_norm
 ones_like = _tree_math.tree_ones_like

optax/tree_utils/__init__.py

@@ -39,6 +39,8 @@
 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_size
 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

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)