Small refactor of gmm in moe.py#3605

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/maxtext:mainfrom

mattdavidow-gmm-clean-a1AI-Hypercomputer/maxtext:mattdavidow-gmm-clean-a1Copy head branch name to clipboard