Fix reductions with side outputs.
Fix reductions with side outputs.

- The current epilogue logic doesn't work for all fusions
- The output indices have to be computed based on a reduction root
- Side outputs have to be written at the end of the loop body

Verified with JAX test suite.
