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

Attention: Patch coverage is 30.00000% with 7 lines in your changes missing coverage. Please review.

Project coverage is 82.72%. Comparing base (744b8be) to head (d7a6ea0).

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.72%   +7.78%     
==========================================
  Files         565      565              
  Lines       55224    55226       +2     
  Branches     8610     8610              
==========================================
+ Hits        41386    45687    +4301     
+ Misses      11880     7427    -4453     
- Partials     1958     2112     +154     

Flag	Coverage Δ
keras	82.53% <30.00%> (+7.76%)	⬆️
keras-jax	63.33% <30.00%> (+<0.01%)	⬆️
keras-numpy	58.53% <10.00%> (?)
keras-openvino	33.79% <10.00%> (?)
keras-tensorflow	63.74% <10.00%> (-0.01%)	⬇️
keras-torch	63.37% <10.00%> (-0.01%)	⬇️

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.