[ROCm]: fix: reduce pipeline temp memory — replace ppermute collectives with lax.slice/pad (PR2) #4192

cj401-amd wants to merge 3 commits into

AI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom

cj401-amd:cj/tmem-fixes-clean-2-pipeline-tmemcj401-amd/maxtext:cj/tmem-fixes-clean-2-pipeline-tmemCopy head branch name to clipboard