[Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support - #5152

[Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support#5152

Rohan-Bierneni wants to merge 1 commit into

rbierneni-gdn-1-ci-hygieneAI-Hypercomputer/maxtext:rbierneni-gdn-1-ci-hygienefrom

rbierneni-gdn-2-fwd-kernelAI-Hypercomputer/maxtext:rbierneni-gdn-2-fwd-kernelCopy head branch name to clipboard