Flash attention support.

[hazemessamm]

I added support for flash attention for PyTorch.

Let me know what do you think about this current implementation so I can add support for JAX and maybe will try for TF.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 8.00000% with 23 lines in your changes missing coverage. Please review.

Project coverage is 79.31%. Comparing base (f4a4725) to head (7e99f06).
Report is 1 commits behind head on master.

Files	Patch %	Lines
keras/src/backend/torch/nn.py	8.00%	23 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20152      +/-   ##
==========================================
- Coverage   79.35%   79.31%   -0.04%     
==========================================
  Files         501      501              
  Lines       47311    47336      +25     
  Branches     8689     8695       +6     
==========================================
+ Hits        37542    37544       +2     
- Misses       8014     8037      +23     
  Partials     1755     1755              

Flag	Coverage Δ
keras	79.16% <8.00%> (-0.04%)	⬇️
keras-jax	62.44% <8.00%> (-0.03%)	⬇️
keras-numpy	57.55% <8.00%> (-0.03%)	⬇️
keras-tensorflow	63.83% <8.00%> (-0.03%)	⬇️
keras-torch	62.47% <8.00%> (-0.03%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the PR -- the code looks good! Please add a unit test.

For the JAX version, I think we'd want to rely on a Pallas kernel. We can get help from the JAX team.

[github-actions]

This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.