Fix `is_causal` fail with compile
Author: Cyrilvallez
Pull request: yes

# What does this PR do?

As per the title. There is a weird issue when running the following:

```py
import torch
from transformers import AutoModelForCausalLM, LlamaConfig

config = LlamaConfig()
config.num_hidden_layers = 2
model = AutoModelForCausalLM.from_config(config).to(0)


model.compile(fullgraph=True)

with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 200), device=0)
    out = model(input_ids)

# Change shape of input
with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 100), device=0)
    out = model(input_ids)
```

It fails with 

```
TypeError: scaled_dot_product_attention(): argument 'is_causal' must be bool, not SymBool
...
RuntimeError: Failed running call_function <built-in function scaled_dot_product_attention>(*(FakeTensor(..., device='cuda:0', size=(1, 32, s0, 128)), FakeTensor(..., device='cuda:0', size=(1, 32, s0, 128)), FakeTensor(..., device='cuda:0', size=(1, 32, s0, 128))), **{'attn_mask': None, 'dropout_p': 0.0, 'scale': 0.08838834764831845, 'is_causal': s0 > 1}):
scaled_dot_product_attention(): argument 'is_causal' must be bool, not SymBool
```

i.e. it somehow traces the `is_causal` assignment as a SymBool instead of evaluating it.

If we pass a mask every time it works, however, as soon as we remove it (which should be equivalent in term of logic) it fails again:

```py
import torch
from transformers import AutoModelForCausalLM, LlamaConfig

config = LlamaConfig()
config.num_hidden_layers = 2
model = AutoModelForCausalLM.from_config(config).to(0)


model.compile(fullgraph=True)

with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 200), device=0)
    out = model(input_ids, attention_mask=torch.ones_like(input_ids))

# Change shape of input but keep a mask
with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 100), device=0)
    out = model(input_ids, attention_mask=torch.ones_like(input_ids))

# Change shape of input but remove mask -> it fails here
with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 100), device=0)
    out = model(input_ids)
```

Simply switching the order of the check in `sdpa_attention` fixes the issue, and ensures that it works independently of passing a mask or not. Not entirely sure why, but it looks like `dynamo` somehow does not treat correctly the shape check if it is done afterwards.

With the current fix, both of the following scenarios work correctly:

Suddenly removing the mask when changing shape:
```py
import torch
from transformers import AutoModelForCausalLM, LlamaConfig

config = LlamaConfig()
config.num_hidden_layers = 2
model = AutoModelForCausalLM.from_config(config).to(0)


model.compile(fullgraph=True)

with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 200), device=0)
    out = model(input_ids, attention_mask=torch.ones_like(input_ids))

with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 100), device=0)
    out = model(input_ids, attention_mask=torch.ones_like(input_ids))

# Stop passing the mask
with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 100), device=0)
    out = model(input_ids)
```

Adding a mask when changing shape:
```py
import torch
from transformers import AutoModelForCausalLM, LlamaConfig

config = LlamaConfig()
config.num_hidden_layers = 2
model = AutoModelForCausalLM.from_config(config).to(0)


model.compile(fullgraph=True)

with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 200), device=0)
    out = model(input_ids)

# Change shape of input
with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 100), device=0)
    out = model(input_ids)

# Add the mask to the inputs
with torch.no_grad():
    input_ids = torch.randint(0, 100, (1, 100), device=0)
    out = model(input_ids, attention_mask=torch.ones_like(input_ids))
```

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_36374). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.