fix: add missing rngs

mesakhcienet · mesakhcienet · commit a1104ade5f15 · 2025-08-19T17:40:46.000+08:00
diff --git a/MaxText/layers/decoders.py b/MaxText/layers/decoders.py
@@ -129,6 +129,7 @@ def __init__(self, config: Config, mesh: Mesh, quant: Optional[Quant] = None, mo
         compute_axis_order=tuple(map(int, self.config.compute_axis_order.split(","))),
         reshape_q=self.config.reshape_q,
         model_mode=model_mode,
+        rngs=self.rngs
     )
 
   def __call__(

Original file line number	Diff line number	Diff line change
`@@ -129,6 +129,7 @@ def __init__(self, config: Config, mesh: Mesh, quant: Optional[Quant] = None, mo`
`129`	`129`	`compute_axis_order=tuple(map(int, self.config.compute_axis_order.split(","))),`
`130`	`130`	`reshape_q=self.config.reshape_q,`
`131`	`131`	`model_mode=model_mode,`
	`132`	`+ rngs=self.rngs`
`132`	`133`	`)`
`133`	`134`
`134`	`135`	`def __call__(`