Fix psum & psum_scatter sharding logic#23

utkarshsharma1 wants to merge 2 commits intomainAI-Hypercomputer/accelerator-microbenchmarks:mainfrom

sharding_logic_fixAI-Hypercomputer/accelerator-microbenchmarks:sharding_logic_fixCopy head branch name to clipboard