Optimize tree_sum compile time using tree_reduce_associative

Changed tree_sum implementation to use jax.tree_util.tree_reduce_associative
when available (JAX >= 0.6.0). Since addition is an associative operation,
tree_reduce_associative can provide better compilation performance.

Testing shows runtime is very close but compile time is significantly
lower (18s vs 23s in reported cases).

For compatibility with older JAX versions (< 0.6.0), the implementation
falls back to the original jax.tree.reduce when tree_reduce_associative
is not available.

Author: Aaryan-549

optax/tree_utils/_tree_math.py

@@ -20,6 +20,7 @@
 
 import jax
 import jax.numpy as jnp
+from jax import tree_util as jtu
 from optax._src import numerics
 
 
@@ -176,7 +177,12 @@ def tree_sum(tree: Any) -> jax.typing.ArrayLike:
     a scalar value.
   """
   sums = jax.tree.map(jnp.sum, tree)
-  return jax.tree.reduce(operator.add, sums, initializer=0)
+  # Use tree_reduce_associative for better compile time performance when
+  # available (JAX >= 0.6.0), otherwise fall back to tree.reduce.
+  if hasattr(jtu, 'tree_reduce_associative'):
+    return jtu.tree_reduce_associative(operator.add, sums)
+  else:
+    return jax.tree.reduce(operator.add, sums, initializer=0)
 
 
 def tree_max(tree: Any) -> jax.typing.ArrayLike: