fixing var order

parambole · parambole · commit c22ceebc2387 · 2025-06-17T02:24:33.000Z
diff --git a/MaxText/tests/multi_token_prediction_test.py b/MaxText/tests/multi_token_prediction_test.py
@@ -125,11 +125,10 @@ class MTPBlockTestModel(nn.Module):
 
   def setup(self):
 
-    self.output_head = models.OutputHead(config=self.config, shared_embedding=self.shared_embedding)
-
     self.shared_embedding = embeddings.Embed(
         num_embeddings=self.config.vocab_size, features=self.config.base_emb_dim, name="token_embedder", config=self.config
     )
+    self.output_head = models.OutputHead(config=self.config, shared_embedding=self.shared_embedding)
     self.mtp_block = multi_token_prediction.MultiTokenPredictionBlock(
         config=self.config, mesh=self.mesh, name="mtp_block", transformer_layer_module=blocks.DecoderLayer
     )

Original file line number	Diff line number	Diff line change
`@@ -125,11 +125,10 @@ class MTPBlockTestModel(nn.Module):`
`125`	`125`
`126`	`126`	`def setup(self):`
`127`	`127`
`128`		`- self.output_head = models.OutputHead(config=self.config, shared_embedding=self.shared_embedding)`
`129`		`-`
`130`	`128`	`self.shared_embedding = embeddings.Embed(`
`131`	`129`	`num_embeddings=self.config.vocab_size, features=self.config.base_emb_dim, name="token_embedder", config=self.config`
`132`	`130`	`)`
	`131`	`+ self.output_head = models.OutputHead(config=self.config, shared_embedding=self.shared_embedding)`
`133`	`132`	`self.mtp_block = multi_token_prediction.MultiTokenPredictionBlock(`
`134`	`133`	`config=self.config, mesh=self.mesh, name="mtp_block", transformer_layer_module=blocks.DecoderLayer`
`135`	`134`	`)`