It worth not hiding exceptions

[WIgor]

What do you think about logging of exceptions, instead of hiding them? + There is an example of such hidden exception when head_shards and q_seq_shards were just not initialized because get_dist()==None.

And despite hardware supports flash attention - this code silently falls back to the standard with O(N^2) memory.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 30.00000% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.74%. Comparing base (744b8be) to head (fa66185).
⚠️ Report is 65 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/nn.py	30.00%	6 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21432      +/-   ##
==========================================
+ Coverage   74.94%   82.74%   +7.79%     
==========================================
  Files         565      567       +2     
  Lines       55224    56466    +1242     
  Branches     8610     8825     +215     
==========================================
+ Hits        41386    46720    +5334     
+ Misses      11880     7581    -4299     
- Partials     1958     2165     +207     

Flag	Coverage Δ
keras	82.54% <30.00%> (+7.77%)	⬆️
keras-jax	63.81% <30.00%> (+0.47%)	⬆️
keras-numpy	58.30% <10.00%> (?)
keras-openvino	34.63% <10.00%> (?)
keras-tensorflow	64.25% <10.00%> (+0.50%)	⬆️
keras-torch	63.87% <10.00%> (+0.49%)	⬆️

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.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request enhances the robustness and observability of the dot_product_attention function by adding logging for fallback scenarios and refactoring the attention implementation selection. The changes improve the user's ability to diagnose performance and compatibility issues.

[divyashreepathihalli]

Thanks for the PR! Please run pre-commit run --all-files

[hertschuh]

@WIgor ,

Thanks for the PR! Can you reformat the code? Then I can merge it.

[WIgor]

As far as I see it does nothing

~/Projects/keras on  patch-1 ⌚ 6:07:26
$ pre-commit run --all-files
api_gen..................................................................Passed
ruff.....................................................................Passed
ruff-format..............................................................Passed
(ml)
~/Projects/keras on  patch-1 ⌚ 6:08:19
$ git status
On branch patch-1
nothing to commit, working tree clean
(ml)

[WIgor]

Ahh ok - it seems working now.