Fix: Remove deprecated .path access in Muon optimizer for TF 2.16+ compatibility #21797
Conversation
…andling and tests
Summary of ChangesHello @utsab345, 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 compatibility issue within the Keras Muon optimizer, specifically targeting an AttributeError that arises in TensorFlow 2.16+ due to the removal of the .path attribute from ResourceVariable objects. The core change involves updating the optimizer's internal logic to use the .name attribute for variable identification, ensuring seamless operation with modern TensorFlow versions. Additionally, the PR enhances the _should_use_adamw method with improved variable dimension checks and more resilient handling of layer exclusion patterns. The changes are thoroughly validated with a suite of new, targeted unit tests, and a minor refinement to image saving utilities is also included. Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request effectively addresses the AttributeError for variable.path in recent TensorFlow versions by migrating to variable.name in the Muon optimizer. The changes are well-implemented and include several improvements:
- The logic in
_should_use_adamwis corrected to properly handle 4D tensors as intended for the Muon update step. - Robustness is improved by adding error handling for invalid regex patterns in
exclude_layers. - The test suite for the Muon optimizer has been significantly improved by refactoring a large test into smaller, more focused tests with clear purposes, and by adding a specific test to prevent regressions of the
.pathattribute issue.
The changes look solid. I have one suggestion to further improve the robustness of the error handling in _should_use_adamw.
| except re.error: | ||
| # Skip invalid regex patterns in exclude_layers | ||
| continue |
There was a problem hiding this comment.
The try...except block is a great addition for robustness against invalid regex patterns. However, re.search can also raise a TypeError if an element in self.exclude_layers is not a string or a compiled regex pattern (e.g., a number). To make this even more robust, consider catching TypeError as well.
| except re.error: | |
| # Skip invalid regex patterns in exclude_layers | |
| continue | |
| except (re.error, TypeError): | |
| # Skip invalid regex patterns in exclude_layers | |
| continue |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #21797 +/- ##
==========================================
+ Coverage 82.63% 82.65% +0.01%
==========================================
Files 577 577
Lines 59321 59331 +10
Branches 9300 9303 +3
==========================================
+ Hits 49020 49040 +20
+ Misses 7913 7890 -23
- Partials 2388 2401 +13
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:
|
utsab345
force-pushed
the
fix-muon-variable-path-issue-21793
branch
from
586b5c0 to
4a299d3
Compare
|
|
||
| def _muon_update_step(self, gradient, variable, lr): | ||
| m = self.adam_momentums[variable.path] | ||
| m = self.adam_momentums[variable.name] |
There was a problem hiding this comment.
This issue is that variable names are not unique. For instance, all Dense layers have a "kernel" variable. So multiple variables will clobber one another.
4 participants
Summary
Fixes keras.optimizers.Muon failing with AttributeError: 'ResourceVariable' object has no attribute 'path' in Keras 3 / TF 2.16–2.20.
Changes
.pathreferences with.namefor variable identification._should_use_adamw()logic to match modern Keras 3 variable handling..pathattribute accessResult
All 11 tests passed locally on TensorFlow 2.16+ (Windows).
Closes #21793