Fix TF int64 promotion issue. #21679
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @james77777778, I'm Gemini Code Assist1! 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 addresses a critical issue in the TensorFlow backend where int64 data types were being incorrectly remapped to int32, causing problems with GPU operations. The fix involves selectively allowing int64 to persist for TensorFlow, aligning with its requirements, while also adjusting test suites to validate this behavior and temporarily bypass related uint32 promotion complexities.
Highlights
- TensorFlow int64 dtype preservation: Modified the
BIT64_TO_BIT32_DTYPEmapping to conditionally retainint64for the TensorFlow backend, preventing its conversion toint32to support GPU operations. - Test adjustments for TensorFlow: Introduced a new test case to verify
int64promotion behavior specifically for TensorFlow and skipped tests involvinguint32for the TensorFlow backend due to ongoing promotion complexities. - Removed BIT64_TO_BIT16_DTYPE: The
BIT64_TO_BIT16_DTYPEmapping dictionary was removed fromkeras/src/backend/common/dtypes.py.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request addresses an issue with int64 dtype promotion on the TensorFlow backend. The change correctly modifies the dtype remapping to preserve int64 for TensorFlow, which is required for certain GPU operations. The addition of a specific test case for int64 promotion on TensorFlow is a great way to ensure this fix is effective and prevents future regressions. The decision to skip tests that fail due to the intended divergence from JAX's behavior is a reasonable, pragmatic approach. I have one suggestion to improve the new test's structure for better consistency.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #21679 +/- ##
=======================================
Coverage 82.56% 82.56%
=======================================
Files 572 572
Lines 58066 58065 -1
Branches 9089 9089
=======================================
+ Hits 47940 47941 +1
+ Misses 7805 7804 -1
+ Partials 2321 2320 -1
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| # TODO(hongyu): Re-enable uint32 tests when we figure out how to handle | ||
| # (uint32, int*) -> int64 promotion. JAX doesn't natively support int64. |
There was a problem hiding this comment.
It may be helpful to explain the TODO in a little more detail, so that we can provide enough context to other contributors who may want to work on this.
There was a problem hiding this comment.
Updated.
# TODO(hongyu): Re-enable uint32 tests once we determine how to handle
# dtypes.result_type(uint32, int*) -> int64 promotion.
# Since TF variables require int64 to be placed on the GPU, we
# exclusively enable the int64 dtype for TF. However, JAX does not
# natively support int64, which prevents us from comparing the dtypes.
google-ml-butler
bot
added
kokoro:force-run
ready to pull
google-ml-butler
bot
removed
the
ready to pull
google-ml-butler
bot
added
kokoro:force-run
ready to pull
Fix #21677
The root cause of the issue is that #21604 forced 64bit dtypes to be remapped to 32bit dtypes to be consistent with JAX's dtype promotion rules. However, TF requires int64 variables for GPU ops.
This PR fixes the issue by exclusively loosening the constraint for the TF backend and skipping the tests that were causing the errors.
cc @JyotinderSingh