[ROCm]: fix: JAX/TE sharding compatibility and tmem reduction foundations (PR1)#4191

cj401-amd wants to merge 5 commits into

AI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom

cj401-amd:cj/tmem-fixes-clean-1-jax-sharding-compatcj401-amd/maxtext:cj/tmem-fixes-clean-1-jax-sharding-compatCopy head branch name to clipboard