Should `tree_sum` use `jax.tree.reduce_associative`?

[SobhanMP]

In tree_sum is implemented using jax.tree.map and jax.tree.reduce but shouldn't jax.tree.reduce_associative be used instead? The runtime seems to be very close but the compile time is significantly lower (18s vs 23s) on the one problem I tested.

[Aaryan-549]

Hi, I have opened a pull request #1503 that addresses this issue.

I implemented the change to use jax.tree_util.tree_reduce_associative as suggested, which should help optimize the compile time. Please take a look when you have a chance!

[rdyro]

Can you point to your use case where the compile time is reduced?

Looking at tree_reduce_associative it uses a different order of summation (associative vs cumulative from left to right), but it is unclear to me how this interacts with the compiler.

Given this is an often-used optax function and this change might affect numerics, I'd prefer to not make any changes to it unless we can identify clear benefits.

@SobhanMP @Aaryan-549

[Aaryan-549]

Hey @rdyro,
The primary benefit here is for PyTrees with a very large number of leaves (e.g., large flat models or optimizers handling many separate state tensors).

Standard tree_reduce performs a sequential fold, creating a computation graph with linear depth O(N) relative to the number of leaves. For trees with thousands of leaves, this results in deep dependency chains that can significantly slow down XLA compilation. tree_reduce_associative performs the reduction in a tree-structure, reducing the graph depth to O(log N).

Regarding numerics: You are correct that this changes the summation order, which will result in slightly different floating-point values compared to a left-to-right accumulation. However, this associative reduction pattern is standard for parallel reductions (like jax.lax.psum) and is usually considered an acceptable trade-off for the compilation speedup in large-scale settings.

If you think the numerical drift is too risky for tree_sum specifically in the context of Optax, I am happy to close this, but I believe the compile time gains for large model users could be significant.

[rdyro]

The graph depth argument sounds reasonable, but we'd need to check with XLA engineers whether it actually slows down compilation.

Because of a change in numerics, I'm against changing the default, but I'd be happy introduce an option to use the associative reduction.

[Aaryan-549]

Sounds great!
I will update the PR to add an optional argument (e.g., acc_dtype or just associative=False) to tree_sum so users can opt-in to the faster compilation path if they need it.