Fix ConcretizationTypeError in dot_product_attention with traced masks

[Abhinavexists]

Fix: #21808

using dot_product_attention with a mask inside JAX control flow operations (e.g., while_loop, scan) on TPU, the code raised a ConcretizationTypeError

Added a traced mask detection check before attempting to use TPU flash attention

[gemini-code-assist]

Summary of Changes

Hello @Abhinavexists, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a specific ConcretizationTypeError that arose when dynamic attention masks, represented as JAX Tracer objects, were passed to the dot_product_attention function within JAX's control flow on TPU devices. The fix ensures that the system gracefully handles these traced masks by bypassing the incompatible TPU flash attention mechanism, thereby improving the robustness and compatibility of attention operations within JAX's tracing environment.

Highlights

• Error Resolution: Addresses a ConcretizationTypeError that occurred when dot_product_attention was used with traced masks within JAX control flow operations on TPUs.
• Flash Attention Logic: Introduces a check within the dot_product_attention function to detect if the provided mask is a JAX Tracer and, if so, disables the use of TPU flash attention to prevent the error.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[codecov-commenter]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.57%. Comparing base (edbf8f5) to head (ffc4bd6).
⚠️ Report is 9 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/nn.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21853      +/-   ##
==========================================
+ Coverage   82.47%   82.57%   +0.10%     
==========================================
  Files         577      577              
  Lines       59508    59587      +79     
  Branches     9332     9347      +15     
==========================================
+ Hits        49080    49205     +125     
+ Misses       8015     7976      -39     
+ Partials     2413     2406       -7     

Flag	Coverage Δ
keras	82.39% <0.00%> (+0.09%)	⬆️
keras-jax	62.87% <0.00%> (-0.03%)	⬇️
keras-numpy	57.52% <0.00%> (-0.04%)	⬇️
keras-openvino	34.34% <0.00%> (-0.01%)	⬇️
keras-tensorflow	64.41% <0.00%> (+0.28%)	⬆️
keras-torch	63.58% <0.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.

🚀 New features to boost your workflow:

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

[gemini-code-assist]

Code Review

This pull request fixes a ConcretizationTypeError that occurs when using dot_product_attention with a traced mask on TPUs. The fix involves adding a check to detect if the mask is a JAX tracer and, if so, disabling flash attention. This is a good and direct solution to the problem. I've suggested a small improvement to scope this check specifically to the TPU platform, aligning it more closely with the reported issue and preventing any potential unintended side effects on other platforms like GPU.

[gemini-code-assist]

Code Review

This pull request addresses a ConcretizationTypeError that occurs when using dot_product_attention with a traced mask on TPUs, which typically happens within JAX control flow constructs like scan or while_loop. The fix involves detecting if the attention mask is a JAX Tracer and disabling flash attention in that case to fall back to a compatible implementation. My review focuses on making this fix more precise to avoid potential performance regressions on other hardware platforms. I've suggested restricting the check to only apply on TPUs, aligning it more closely with the described issue.

[hertschuh]

Because on JAX we compile by default, this means that flash_attention is now always disabled. The whole flash attention implementation was for TPU specifically. So I don't think this is the right fix.

Am I misunderstanding something?

[Abhinavexists]

@hertschuh You're absolutely right! Thank you for catching that.

My initial fix using isinstance(mask, jax.core.Tracer) was too broad and would have disabled flash attention even during normal JIT compilation, not just inside control flow operations.

I've revised the fix to instead:

1. Removed the Tracer check
2. Improved the exception handling to specifically catch jax.errors.ConcretizationTypeError

The fix now correctly handles traced masks without disabling flash attention for normal use cases.