flax

implement axes_scan's broadcast semantics via dce

#5540

Merged