Wrap _prepare_4d_causal_attention_mask as a leaf function

[michaelbenayoun]

What does this PR do?

This wraps the _prepare_4d_causal_attention_mask as a FX leaf function for similar reasons than here.

The only consequence it has is that it will not be possible to edit this function by using torch.fx. It is not a big deal at all, but I will remove this constraint as soon as possible.

[amyeroberts]

Change LGTM - thanks for adding!

Happy to merge with @younesbelkada's approval.

cc @patrickvonplaten for reference

[younesbelkada]

The changes look good! My only question is that this seems to wrap _prepare_4d_causal_attention_mask all the time as long as fx is available, is it possible to wrap it only if users perform model tracing?

[michaelbenayoun]

No it needs to happen at the top-module level.

[younesbelkada]

LGTM, reading a bit the docs of the method it just registers the method as a "leaf function" without modifying anything - seems safe to me

https://pytorch.org/docs/stable/_modules/torch/fx/_symbolic_trace.html#wrap

[michaelbenayoun]

It is, the only difference it implies is that it will not be possible to edit this function via torch.fx.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[michaelbenayoun]

Actually I have a workaround that works for my purposes in optimum-neuron, I can reverse the changes or keep them like that, as you wish. It's not a big deal in any case.

[amyeroberts]

I think it's fine to leave as-is - people might want to use torch tracing outside of optimum-neuron.