Support specifying tokamax gmm tile sizes in MaxText#3779

Open

darisoy wants to merge 1 commit intomainfrom

darisoy-gmm-tile

Commits on Apr 29, 2026

Support specifying custom tile sizes for forward and backward passes of Tokamax GMM in MaxText
darisoy
committed