Add missing initializer argument of 0 to tree_reduce in tree_vdot and…

… tree_sum.
google-deepmind · Sep 20, 2024 · dfb10c8 · dfb10c8
1 parent ee63e45
commit dfb10c8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/optax/tree_utils/_tree_math.py b/optax/tree_utils/_tree_math.py
@@ -152,7 +152,7 @@ def tree_vdot(tree_x: Any, tree_y: Any) -> chex.Numeric:
  numerical issues.
  """
  vdots = jtu.tree_map(_vdot_safe, tree_x, tree_y)
- return jtu.tree_reduce(operator.add, vdots)
+ return jtu.tree_reduce(operator.add, vdots, 0)
 
 
 def tree_sum(tree: Any) -> chex.Numeric:
@@ -165,7 +165,7 @@ def tree_sum(tree: Any) -> chex.Numeric:
  a scalar value.
  """
  sums = jtu.tree_map(jnp.sum, tree)
- return jtu.tree_reduce(operator.add, sums)
+ return jtu.tree_reduce(operator.add, sums, 0)
 
 
 def _square(leaf):

diff --git a/optax/tree_utils/_tree_math_test.py b/optax/tree_utils/_tree_math_test.py
@@ -232,6 +232,12 @@ def test_bias_correction_bf16(self):
  tu.tree_bias_correction(m, decay, count),
  custom_message=f'failed with decay={decay}, count={count}')
 
+ def test_empty_tree_reduce(self):
+ # assert False
+ for tree in [{}, (), [], None, {"key": [None, [None]]}]:
+ tu.tree_sum(tree)
+ tu.tree_vdot(tree, tree)
+
 
 if __name__ == '__main__':
  absltest.main()