Refactor to move logic related to allreduce to the method for it.#786

copybara-service[bot] wants to merge 1 commit intomainopenxla/shardy:mainfrom

test_812756921openxla/shardy:test_812756921Copy head branch name to clipboard