Internal change

PiperOrigin-RevId: 542883908

Author: emilyfertig

jaxopt/_src/tree_util.py

@@ -167,8 +167,8 @@ def tree_gram(a):
 
 def tree_inf_norm(tree_x):
   """Computes the infinity norm of a pytree."""
-  leaf_inf_norm = tree_map(lambda x: jnp.max(jnp.abs(x)), tree_x)
-  return tree_reduce(jnp.maximum, leaf_inf_norm)
+  leaves_vec = tree_leaves(tree_map(jnp.ravel, tree_x))
+  return jnp.max(jnp.abs(jnp.concatenate(leaves_vec)))
 
 
 def tree_where(cond, a, b):

tests/tree_util_test.py

@@ -129,6 +129,22 @@ def test_tree_l2_norm(self):
     got = tree_util.tree_l2_norm(self.tree_A)
     self.assertAllClose(expected, got)
 
+  def test_tree_inf_norm(self):
+    expected = jnp.max(jnp.abs(self.array_A))
+    got = tree_util.tree_inf_norm(self.array_A)
+    self.assertAllClose(expected, got)
+
+    expected = jnp.maximum(jnp.max(jnp.abs(self.tree_A[0])),
+                           jnp.max(jnp.abs(self.tree_A[1])))
+    got = tree_util.tree_inf_norm(self.tree_A)
+    self.assertAllClose(expected, got)
+
+    tree_a_with_empty_leaves = ({'a': self.tree_A, 'b': jnp.array([])},
+                                jnp.array([]))
+    expected = tree_util.tree_inf_norm(self.tree_A)
+    got = tree_util.tree_inf_norm(tree_a_with_empty_leaves)
+    self.assertAllClose(expected, got)
+
   def test_tree_zeros_like(self):
     tree = tree_util.tree_zeros_like(self.tree_A)
     self.assertAllClose(tree_util.tree_l2_norm(tree), 0.0)