Repository: huggingface/transformers
Pull request: Fix lifted_tensor in Gemma3n export which dynamo can't reason about

When using the torch.export path with this model we hit an issue on correction_coefs being modified with unknown side-effects the dynamo can't catch. So it lifts this tensor unnecessarily, and fails assuming mutation is needed RuntimeError: Constant model.model.language_model.layers.0.lifted_tensor_4 is mutated in the forward method. Pls register it as buffer.

This change makes the subsequent use recognisable as an out variant to prevent this.

As this is a minor adjustment with no functionality added I've just checked that the basic flow and gemma3n tests pass as in the guide:
```
pip install -e ".[dev]"
python utils/modular_model_converter.py gemma3n
RUN_SLOW=1 pytest tests/models/gemma3n/
make style
make check-repo
```

Hope this needs attn of:
- multimodal models: @zucchini-nlp

Changed files:
filename: src/transformers/models/gemma3n/modeling_gemma3n.py
patch: @@ -1145,12 +1145,14 @@ def correct(self, predictions: torch.Tensor, activated: torch.Tensor) -> torch.T
         innovation = innovation.repeat(self.config.altup_num_inputs, 1, 1, 1)  # Repeat on dim0 to match predictions
 
         if self.training and self.config.altup_coef_clip is not None:
-            self.correction_coefs.weight.data.clamp_(-self.config.altup_coef_clip, self.config.altup_coef_clip)
+            weight = self.correction_coefs.weight.clamp(-self.config.altup_coef_clip, self.config.altup_coef_clip)
+            all_coefs = torch.nn.functional.linear(modalities, weight, bias=None) + 1.0
+        else:
+            all_coefs = self.correction_coefs(modalities) + 1.0
 
         # all_coefs adapted from jax.numpy.einsum("...p,pi->...i", ...)
         # Permute to (altup_num_inputs, batch_size, num_tokens) as the last dim is a scalar applied to each altup input
         # and expand on dim1 for broadcastability
-        all_coefs: torch.Tensor = self.correction_coefs(modalities) + 1.0
         all_coefs = all_coefs.permute(2, 0, 1).unsqueeze(-1)
 
         corrected = torch.mul(innovation, all_coefs)

filename: src/transformers/models/gemma3n/modular_gemma3n.py
patch: @@ -1694,12 +1694,14 @@ def correct(self, predictions: torch.Tensor, activated: torch.Tensor) -> torch.T
         innovation = innovation.repeat(self.config.altup_num_inputs, 1, 1, 1)  # Repeat on dim0 to match predictions
 
         if self.training and self.config.altup_coef_clip is not None:
-            self.correction_coefs.weight.data.clamp_(-self.config.altup_coef_clip, self.config.altup_coef_clip)
+            weight = self.correction_coefs.weight.clamp(-self.config.altup_coef_clip, self.config.altup_coef_clip)
+            all_coefs = torch.nn.functional.linear(modalities, weight, bias=None) + 1.0
+        else:
+            all_coefs = self.correction_coefs(modalities) + 1.0
 
         # all_coefs adapted from jax.numpy.einsum("...p,pi->...i", ...)
         # Permute to (altup_num_inputs, batch_size, num_tokens) as the last dim is a scalar applied to each altup input
         # and expand on dim1 for broadcastability
-        all_coefs: torch.Tensor = self.correction_coefs(modalities) + 1.0
         all_coefs = all_coefs.permute(2, 0, 1).unsqueeze(-1)
 
         corrected = torch.mul(innovation, all_coefs)