Add cuda profiler hook to paxml#1089

shawnwang18 wants to merge 1 commit intoNVIDIA:mainNVIDIA/JAX-Toolbox:mainfrom

shawnwang18:shawnw/add_paxml_patchshawnwang18/JAX-Toolbox:shawnw/add_paxml_patchCopy head branch name to clipboard