Fixed issue with dot_product_attention when using TPU. #21254

pctablet505 · 2025-05-06T06:34:02Z

No description provided.

codecov-commenter · 2025-05-06T06:45:18Z

Codecov Report

Attention: Patch coverage is 14.51613% with 53 lines in your changes missing coverage. Please review.

Project coverage is 82.53%. Comparing base (6ddaefb) to head (f60811e).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/nn.py	14.51%	48 Missing and 5 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21254      +/-   ##
==========================================
- Coverage   82.59%   82.53%   -0.07%     
==========================================
  Files         564      564              
  Lines       54594    54642      +48     
  Branches     8483     8495      +12     
==========================================
+ Hits        45092    45098       +6     
- Misses       7415     7454      +39     
- Partials     2087     2090       +3

Flag	Coverage Δ
keras	`82.34% <14.51%> (-0.07%)`	⬇️
keras-jax	`63.62% <14.51%> (-0.05%)`	⬇️
keras-numpy	`58.74% <0.00%> (-0.06%)`	⬇️
keras-openvino	`32.96% <0.00%> (-0.03%)`	⬇️
keras-tensorflow	`64.03% <0.00%> (-0.06%)`	⬇️
keras-torch	`63.69% <0.00%> (-0.06%)`	⬇️

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.

fchollet

Thanks for the PR!

fchollet · 2025-05-07T04:52:02Z

keras/src/backend/jax/nn.py

        strides: a sequence of `N` integers, representing the inter-window
-            strides (default: `(1, ..., 1)`).
+        strides (default: `(1, ..., 1)`).


Please add indent

fchollet · 2025-05-07T04:52:09Z

keras/src/backend/jax/nn.py

-            "Sharding along sequence dimension not allowed in tpu kernel "
-            "attention"
+            "Sharding along sequence dimension not allowed"
+            " in tpu kernel attention"


fchollet · 2025-05-07T04:52:18Z

keras/src/backend/jax/nn.py

+
+    Args:
+        query: Queries with shape `[batch, time, heads,
+          depth_k]`.


Please use 4-space indent

Corrected indentation in doc string

fchollet

LGTM, thank you

…-team#21254)" This reverts commit d8f3f70.

…" (#21329) This reverts commit d8f3f70.

…s-team#21254)" (keras-team#21329) This reverts commit 81821e0.

…after addressing cuDNN/FlashAttention API updates (#21333) * Update nn.py * Update nn.py * Update nn.py * Update nn.py * Update nn.py Corrected indentation in doc string * Update nn.py * Update random_grayscale.py Fixed issue with passing a single image without batch dimension. * Update keras/src/layers/preprocessing/image_preprocessing/random_grayscale.py Co-authored-by: Jyotinder Singh <[email protected]> * Update random_grayscale_test.py Test case for unbatched inputs * code reformat * Update random_grayscale_test.py Testcase for checking both unbatched and batched single image inputs. * changed compute_output_spec There was a bug, and it was causing cycle in graph. * Update random_grayscale.py removed the use of tree.map_structure * Reapply "Fixed issue with dot_product_attention when using TPU. (#21254)" (#21329) This reverts commit 81821e0. * Improve error handling in _can_use_flash_attention for better debugging Enhanced the _can_use_flash_attention function to provide more detailed error messages when flash attention compatibility checks fail. Changes: - Replace generic exception catching with specific error propagation - When raise_error=True, directly re-raise original exceptions from check_layout() and check_is_flash_attention() functions - Preserve detailed error context from JAX internal validation functions - Maintain existing behavior when raise_error=False (returns False) This improves debugging experience by surfacing specific technical details about tensor layout incompatibilities, cuDNN version requirements, and other flash attention compatibility issues. Relates to keras-hub PR #2257 and addresses flash attention debugging needs. * Revert "Improve error handling in _can_use_flash_attention for better debugging" This reverts commit 7a0c547. * Fix JAX API compatibility and improve error handling in `_can_use_flash_attention` Changes: - Add missing q_offsets=None and kv_offsets=None parameters to check_layout() call to match updated JAX function signature - Replace bare `except:` with `except Exception as e:` and `raise e` to preserve detailed error messages from JAX validation functions - Maintain existing fallback behavior when raise_error=False This resolves compatibility issues with newer JAX versions and improves debugging experience by surfacing specific technical details about flash attention compatibility failures. * Updated `dot_product_attention` Simplified the check for `flasth_attention` by removing redundant checks that are already done in `_can_use_flash_attention`. * Update nn.py * Update nn.py --------- Co-authored-by: Jyotinder Singh <[email protected]>

* Update nn.py * Update nn.py * Update nn.py * Update nn.py * Update nn.py Corrected indentation in doc string * Update nn.py * Update random_grayscale.py Fixed issue with passing a single image without batch dimension. * Update keras/src/layers/preprocessing/image_preprocessing/random_grayscale.py Co-authored-by: Jyotinder Singh <[email protected]> * Update random_grayscale_test.py Test case for unbatched inputs * code reformat * Update random_grayscale_test.py Testcase for checking both unbatched and batched single image inputs. * changed compute_output_spec There was a bug, and it was causing cycle in graph. * Update random_grayscale.py removed the use of tree.map_structure * Reapply "Fixed issue with dot_product_attention when using TPU. (#21254)" (#21329) This reverts commit 81821e0. * Improve error handling in _can_use_flash_attention for better debugging Enhanced the _can_use_flash_attention function to provide more detailed error messages when flash attention compatibility checks fail. Changes: - Replace generic exception catching with specific error propagation - When raise_error=True, directly re-raise original exceptions from check_layout() and check_is_flash_attention() functions - Preserve detailed error context from JAX internal validation functions - Maintain existing behavior when raise_error=False (returns False) This improves debugging experience by surfacing specific technical details about tensor layout incompatibilities, cuDNN version requirements, and other flash attention compatibility issues. Relates to keras-hub PR #2257 and addresses flash attention debugging needs. * Revert "Improve error handling in _can_use_flash_attention for better debugging" This reverts commit 7a0c547. * Fix JAX API compatibility and improve error handling in `_can_use_flash_attention` Changes: - Add missing q_offsets=None and kv_offsets=None parameters to check_layout() call to match updated JAX function signature - Replace bare `except:` with `except Exception as e:` and `raise e` to preserve detailed error messages from JAX validation functions - Maintain existing fallback behavior when raise_error=False This resolves compatibility issues with newer JAX versions and improves debugging experience by surfacing specific technical details about flash attention compatibility failures. * Updated `dot_product_attention` Simplified the check for `flasth_attention` by removing redundant checks that are already done in `_can_use_flash_attention`. * Update nn.py * Update nn.py * Update image.py * Update keras/src/backend/tensorflow/image.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Revert "Update keras/src/backend/tensorflow/image.py" This reverts commit cb7e955. * Update image.py * Update image.py --------- Co-authored-by: Jyotinder Singh <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

pctablet505 added 3 commits May 6, 2025 10:25

Update nn.py

04cd682

Update nn.py

1a74465

Update nn.py

c11eb81

google-ml-butler bot added the size:M label May 6, 2025

google-ml-butler bot assigned gbaned May 6, 2025

Update nn.py

c81e18c

pctablet505 mentioned this pull request May 6, 2025

Update nn.py #21250

Closed

pctablet505 requested a review from divyashreepathihalli May 7, 2025 04:20

google-ml-butler bot added the awaiting review label May 7, 2025

fchollet reviewed May 7, 2025

View reviewed changes

pctablet505 added 2 commits May 7, 2025 10:27

Update nn.py

d938e20

Corrected indentation in doc string

Update nn.py

f60811e

fchollet approved these changes May 7, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels May 7, 2025

fchollet merged commit d8f3f70 into keras-team:master May 7, 2025
7 checks passed

google-ml-butler bot removed awaiting review ready to pull Ready to be merged into the codebase kokoro:force-run labels May 7, 2025

chaosmaster142857 mentioned this pull request May 15, 2025

Splash Attention is Broken on TPU Pods and does not follow keras.config.disable_flash_attention() #21116

Closed

divyashreepathihalli mentioned this pull request May 16, 2025

Update requirements-jax-cuda.txt keras-team/keras-hub#2252

Merged

pctablet505 added a commit to pctablet505/keras that referenced this pull request May 27, 2025

Revert "Fixed issue with dot_product_attention when using TPU. (keras…

36b5dc3

…-team#21254)" This reverts commit d8f3f70.

pctablet505 mentioned this pull request May 27, 2025

Revert "Fixed issue with dot_product_attention when using TPU. " #21329

Merged

fchollet pushed a commit that referenced this pull request May 27, 2025

Revert "Fixed issue with dot_product_attention when using TPU. (#21254)…

81821e0

…" (#21329) This reverts commit d8f3f70.

pctablet505 added a commit to pctablet505/keras that referenced this pull request May 29, 2025

Reapply "Fixed issue with dot_product_attention when using TPU. (kera…

579cc11

…s-team#21254)" (keras-team#21329) This reverts commit 81821e0.

pctablet505 mentioned this pull request May 29, 2025

Re-apply Fixed issue with dot_product_attention when using TPU. #21254 after addressing cuDNN/FlashAttention API updates #21333

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fixed issue with dot_product_attention when using TPU. #21254

Fixed issue with dot_product_attention when using TPU. #21254

Uh oh!

pctablet505 commented May 6, 2025

Uh oh!

codecov-commenter commented May 6, 2025 •

edited

Loading

Uh oh!

fchollet left a comment

Uh oh!

fchollet May 7, 2025

Uh oh!

fchollet May 7, 2025

Uh oh!

fchollet May 7, 2025

Uh oh!

fchollet left a comment

Uh oh!

Uh oh!

Uh oh!

Fixed issue with dot_product_attention when using TPU. #21254

Fixed issue with dot_product_attention when using TPU. #21254

Uh oh!

Conversation

pctablet505 commented May 6, 2025

Uh oh!

codecov-commenter commented May 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

fchollet left a comment

Choose a reason for hiding this comment

Uh oh!

fchollet May 7, 2025

Choose a reason for hiding this comment

Uh oh!

fchollet May 7, 2025

Choose a reason for hiding this comment

Uh oh!

fchollet May 7, 2025

Choose a reason for hiding this comment

Uh oh!

fchollet left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

codecov-commenter commented May 6, 2025 •

edited

Loading