Fix is_causal being a tensor

[IlyasMoutawwakil]

What does this PR do?

We encountered this issue when gpt2 is part of a vision encoder decoder (add_cross_attention=True)

FAILED exporters/onnx/test_exporters_onnx_cli.py::OnnxCLIExportTestCase::test_exporters_cli_pytorch_cpu_639_vision_encoder_decoder_image_to_text_with_past_default_hf_internal_testing_tiny_random_VisionEncoderDecoderModel_vit_gpt2 - TypeError: scaled_dot_product_attention(): argument 'is_causal' must be bool, not Tensor

Usually is_causal is evaluated in is_causal = attention_mask is None and query_states.shape[-2] > 1 and not is_cross_attention by the fact that attention_mask is None is False, skipping the rest, but when attention mask is not None, query_states.shape[-2] > 1 is evaluated and returns a tensor, resulting in sdpa failing because it's expecting a bool.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[ArthurZucker]

Find it a bit weird because I can't reproduce the tensor shape being eval as tensor

>>> torch.ones(*(1,2),1).shape[0] > 1
False

[IlyasMoutawwakil]

Yeah I was surprised at first and thought it might be a behavior change in torch but it's actually related to tracing:

import torch


class model(torch.nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.fc = torch.nn.Linear(1, 1)

    def forward(self, x):
        print("shape > 1: ", x.shape[0] > 1)
        return self.fc(x)


model = model()
x = torch.randn(1, 1)

model(x)  # shape > 1:  False
torch.jit.trace(model, x, check_trace=False)  # shape > 1:  tensor(False)
torch.onnx.export(model, x, "model.onnx")  # shape > 1:  tensor(False)

This should've impacted the gpt2 onnx tests but I guess it wasn't caught because it's only tested as a pure decoder.

[ArthurZucker]

Got it, does the device synch not cause any issue for the exported model afterwards?

[IlyasMoutawwakil]

Got it, does the device synch not cause any issue for the exported model afterwards?

@ArthurZucker didn't see any issues in optimum onnx exported models, but I added a guard for jit tracing just in case.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Okay thanks