Optimize tree_sum compile time using tree_reduce_associative

[Aaryan-549]

In Response to #1498
Changed tree_sum implementation to use jax.tree_util.tree_reduce_associative instead of jax.tree.reduce. 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).

[SobhanMP]

I think this should be
jax.tree.reduce_associative(operator.add, sums, initializer=0)

and the pythonic way would probably use AttributeError and try catch

[Aaryan-549]

Hi @SobhanMP
Thanks for the review! However, jax.tree.reduce_associative doesn't support the initializer parameter - it only accepts (function, tree). That's why the original implementation failed with "Must specify identity for parallel reduction of empty sequence" when encountering empty trees. The current approach with hasattr is appropriate here because:

1. We're checking for API availability across JAX versions (0.5.3 vs 0.6.0+)
2. The empty tree check is necessary since tree_reduce_associative lacks initializer support
3. Using hasattr for version compatibility is a common pattern in the JAX ecosystem

A try/except approach would be:

try:
    leaves = jax.tree.leaves(sums)
    if not leaves:
        return 0
    return jtu.tree_reduce_associative(operator.add, sums)
except AttributeError:
    return jax.tree.reduce(operator.add, sums, initializer=0)

But this doesn't provide much benefit over hasattr for this use case, and the hasattr check makes the version compatibility intent clearer.

[carlosgmartin]

@Aaryan-549 jax.tree.reduce_associative has an identity argument.

[Aaryan-549]

You're absolutely right, thank you for the correction! I've updated the implementation to use identity=0 instead of manually checking for empty trees. Much cleaner now:

if hasattr(jtu, 'tree_reduce_associative'):
    return jtu.tree_reduce_associative(operator.add, sums, identity=0)
else:
    return jax.tree.reduce(operator.add, sums, initializer=0)