From 2a7d783a748896d90ddf2e024e3acf4e445b04f6 Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Wed, 11 Mar 2026 14:14:44 -0700 Subject: [PATCH] remove variable names PiperOrigin-RevId: 882192055 --- src/maxtext/layers/moe.py | 2 +- src/maxtext/models/deepseek_batchsplit.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 70e67490f5..3a54a2a900 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -898,7 +898,7 @@ def gmm( ): tokamax_group_sizes = tokamax.RaggedDotGroupSizes( group_sizes, - representative_value=max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]), + max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]), ) pad_length = self.config.wi_tile_fwd_batch_seq hs_shape = inputs.shape diff --git a/src/maxtext/models/deepseek_batchsplit.py b/src/maxtext/models/deepseek_batchsplit.py index 7cbbe17063..680e17a1d3 100644 --- a/src/maxtext/models/deepseek_batchsplit.py +++ b/src/maxtext/models/deepseek_batchsplit.py @@ -807,7 +807,7 @@ def gmm( tokamax_group_sizes = tokamax.RaggedDotGroupSizes( group_sizes, - representative_value=max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]), + max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]), ) if config.use_qwix_quantization: output = megablox.gmm(