DSV4: Add mhc_split_axis_contraction to contract the mHC rate and embed axes separately - #4912

DSV4: Add mhc_split_axis_contraction to contract the mHC rate and embed axes separately#4912

systalyze-ai wants to merge 1 commit into

AI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom

systalyze-ai:upstream-pr/mhc-split-axis-contractionsystalyze-ai/maxtext:upstream-pr/mhc-split-axis-contractionCopy head branch name to clipboard