Merge pull request #3746 from AI-Hypercomputer:aireen/fix_gemma4_dtype

PiperOrigin-RevId: 905389377

Author: Google-ML-Automation

src/maxtext/models/gemma4.py

@@ -370,7 +370,7 @@ def __call__(
 
     next_layer_addition = mlp_lnx + residual
     layer_output = next_layer_addition
-    layer_output = layer_output * self.layer_scalar.value
+    layer_output = layer_output * jnp.asarray(self.layer_scalar.value, cfg.dtype)
 
     layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names)