can't create torch_jit_trace script for a Mistral model
Author: helloWorld199
Pull request: no

### System Info

transformers version: 4.34.0
torch version: 2.3.0
python version: 3.10.12

### Who can help?

_No response_

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction


```
import torch
from transformers import MistralConfig, MistralModel, MistralForCausalLM

configuration = MistralConfig(hidden_size = 512, intermediate_size = 512*4, num_attention_heads = 16, num_hidden_layers = 6, num_key_value_heads = 4, vocab_size = 700)
mistral_model = MistralForCausalLM(configuration)

mistral_output= model(input_ids=example_input)
traced_script_module = torch.jit.trace(mistral_model, [example_input, mistraloutput.past_key_values])
```

### Output error ###
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-16-bb4ba78240fd>](https://localhost:8080/#) in <cell line: 1>()
----> 1 traced_script_module = torch.jit.trace(mistral_model, [example_input, mistraloutput[1]])

11 frames
[/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py](https://localhost:8080/#) in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_kwarg_inputs, _store_inputs)
    818             else:
    819                 raise RuntimeError("example_kwarg_inputs should be a dict")
--> 820         return trace_module(
    821             func,
    822             {"forward": example_inputs},

[/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py](https://localhost:8080/#) in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_inputs_is_kwarg, _store_inputs)
   1086             else:
   1087                 example_inputs = make_tuple(example_inputs)
-> 1088                 module._c._create_method_from_trace(
   1089                     method_name,
   1090                     func,

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1530             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531         else:
-> 1532             return self._call_impl(*args, **kwargs)
   1533 
   1534     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1539                 or _global_backward_pre_hooks or _global_backward_hooks
   1540                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541             return forward_call(*args, **kwargs)
   1542 
   1543         try:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _slow_forward(self, *input, **kwargs)
   1520                 recording_scopes = False
   1521         try:
-> 1522             result = self.forward(*input, **kwargs)
   1523         finally:
   1524             if recording_scopes:

[/usr/local/lib/python3.10/dist-packages/transformers/models/mistral/modeling_mistral.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1043 
   1044         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1045         outputs = self.model(
   1046             input_ids=input_ids,
   1047             attention_mask=attention_mask,

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1530             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531         else:
-> 1532             return self._call_impl(*args, **kwargs)
   1533 
   1534     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1539                 or _global_backward_pre_hooks or _global_backward_hooks
   1540                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541             return forward_call(*args, **kwargs)
   1542 
   1543         try:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _slow_forward(self, *input, **kwargs)
   1520                 recording_scopes = False
   1521         try:
-> 1522             result = self.forward(*input, **kwargs)
   1523         finally:
   1524             if recording_scopes:

[/usr/local/lib/python3.10/dist-packages/transformers/models/mistral/modeling_mistral.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
    886                 )
    887 
--> 888         attention_mask = self._prepare_decoder_attention_mask(
    889             attention_mask,
    890             (batch_size, seq_length),

[/usr/local/lib/python3.10/dist-packages/transformers/models/mistral/modeling_mistral.py](https://localhost:8080/#) in _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length, sliding_window)
    804         if attention_mask is not None:
    805             # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
--> 806             expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
    807                 inputs_embeds.device
    808             )

[/usr/local/lib/python3.10/dist-packages/transformers/models/mistral/modeling_mistral.py](https://localhost:8080/#) in _expand_mask(mask, dtype, tgt_len)
     99     Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    100     """
--> 101     bsz, src_len = mask.size()
    102     tgt_len = tgt_len if tgt_len is not None else src_len
    103 

AttributeError: 'tuple' object has no attribute 'size'

### Expected behavior

I expect the traced_script_module to be correctly initialized

If I do the same thing with a GPT2 model, as in this code snippet, it works, even if the structure of the past_key_values is the same:
```
from transformers import GPT2LMHeadModel, GPT2Config
gpt2_config = GPT2Config()
gpt2_model = GPT2LMHeadModel(gpt2_config)
gpt2_output = gpt2_model(example_input)
traced_script_module = torch.jit.trace(gpt2_model, [example_input, gpt2_output.past_key_values])
```
What should I do to make the script work for mistral aswell?

Thanks! 

Comments:
cc @ArthurZucker, I don't think the Mistral model is jittable just yet

Hey! The `_prepare_decoder_attention_mask` is the old function, which was removed in favor of https://github.com/huggingface/transformers/blob/cd8d08df4be23aa55873ecb58aa1f646d279107c/src/transformers/models/mistral/modeling_mistral.py#L948-L951

you are probably not using the latest version of transformers! 

`pip install -U transformers` should fix it

oh yep that was it @ArthurZucker, thanks!