Add Transformer Engine support to Paxml#46

ashors1 wants to merge 12 commits into

google:maingoogle/paxml:mainfrom

nvjax-svc-0:patch/te_supportnvjax-svc-0/paxml:patch/te_supportCopy head branch name to clipboard