Fix #21105: JAX Backend - Custom Gradient Variable Capture Issue in keras.ops.custom_gradient #21783
Conversation
- torch-xla is not available for Windows platform - Manually installed tensorflow-cpu, torch, jax, and flax - Fixed protobuf version conflicts (downgraded to <6.0.0) - Tests now run successfully without ModuleNotFoundError
Summary of ChangesHello @Mayankvlog, 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 bug in the Keras JAX backend where 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 correctly fixes a variable capture issue with custom_gradient in the JAX backend. The implementation is clean, and the added tests are thorough, ensuring the issue is resolved and preventing regressions.
However, the pull request also includes changes related to a jit_compile issue for the Torch backend (issue #21647), which is not reflected in the PR's title or description. Combining unrelated fixes can complicate the review process and make the commit history less clear. In the future, please consider creating separate pull requests for distinct issues to improve clarity and maintainability.
keras/src/backend/jax/core.py
Outdated
| def wrapper(*args, **kwargs): | ||
| # Convert Variable objects to their values | ||
| def _convert_arg(arg): | ||
| if isinstance(arg, Variable): | ||
| return arg.value | ||
| return arg | ||
|
|
||
| args = tree.map_structure(_convert_arg, args) | ||
| kwargs = tree.map_structure(_convert_arg, kwargs) | ||
| return fun(*args, **kwargs) | ||
|
|
||
| return jax.custom_gradient(fun=wrapper) |
There was a problem hiding this comment.
The current implementation redefines the _convert_arg function on every call to the decorated function, which is inefficient. This helper function can be defined once outside the wrapper to avoid this overhead. Additionally, renaming it to _convert_variable_to_value would make its purpose clearer, following the style guide's preference for descriptive names.1
| def wrapper(*args, **kwargs): | |
| # Convert Variable objects to their values | |
| def _convert_arg(arg): | |
| if isinstance(arg, Variable): | |
| return arg.value | |
| return arg | |
| args = tree.map_structure(_convert_arg, args) | |
| kwargs = tree.map_structure(_convert_arg, kwargs) | |
| return fun(*args, **kwargs) | |
| return jax.custom_gradient(fun=wrapper) | |
| def _convert_variable_to_value(arg): | |
| if isinstance(arg, Variable): | |
| return arg.value | |
| return arg | |
| def wrapper(*args, **kwargs): | |
| # Convert Variable objects to their values | |
| args = tree.map_structure(_convert_variable_to_value, args) | |
| kwargs = tree.map_structure(_convert_variable_to_value, kwargs) | |
| return fun(*args, **kwargs) | |
| return jax.custom_gradient(fun=wrapper) |
Style Guide References
Footnotes
-
Argument names should be intuitive and easy to remember, and their meaning should be clear from the name. Overly generic names should be avoided. ↩
…ng errors - Fixed custom_gradient in JAX backend to extract Variable values automatically - Improved code structure by moving helper function outside wrapper - Fixed EfficientNetV2B2 import to use direct module import - Fixed all Ruff linting errors (line length, unused imports/variables) - Tests now pass without requiring manual .value access on Variables
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #21783 +/- ##
===========================================
- Coverage 82.63% 63.78% -18.86%
===========================================
Files 577 578 +1
Lines 59316 59404 +88
Branches 9300 9313 +13
===========================================
- Hits 49018 37893 -11125
- Misses 7910 19081 +11171
- Partials 2388 2430 +42
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:
|
- Changed input size from 64x64 to 224x224 (minimum supported by EfficientNetV2) - Fixed EfficientNetV2B0 import to use direct module path - Resolves ValueError: Input size must be at least 32x32 - Resolves ImportError for EfficientNetV2B0
…input_shape validation This commit addresses three issues that were causing CI failures: 1. Fixed JAX Backend custom_gradient with Variables (Issue keras-team#21105) - Problem: Variables passed to custom_gradient in JAX backend caused 'TypeError: NoneType object is not callable' - Root cause: JAX copies Variables during tracing, causing both _value and _initializer to become None - Solution: * Modified custom_gradient wrapper to properly convert Variables to values * Added fallback in __jax_array__ to handle uninitialized Variables - Added test: test_custom_gradient_with_variable in keras/src/ops/core_test.py 2. Fixed obtain_input_shape validation for channels_first format - Problem: Confusing error when users provide input_shape in wrong format (e.g., (224,224,3) when (3,224,224) expected for channels_first) - Solution: Added validation to detect format mismatch with clear error message - Updated efficientnet_v2_jit_test.py to use correct channels_first format 3. Code format fixes - Fixed line length violations - Fixed import ordering - Removed unused imports Files modified: - keras/src/backend/jax/core.py - keras/src/ops/core_test.py - keras/src/applications/imagenet_utils.py - keras/src/applications/efficientnet_v2_jit_test.py - test_custom_gradient_jax_variable.py All tests passing with JAX backend.
3 participants
Resolved issue #21105 where keras.ops.custom_gradient incorrectly captured Variable objects instead of their values when using the JAX backend. This caused gradient computation failures in custom quantization layers. The problem occurred when passing keras.Variable instances directly to functions decorated with @ops.custom_gradient, where the gradient function (roundpass_grad) captured the variable object rather than its underlying value.
Root Cause: JAX backend's custom gradient implementation didn't properly extract values from Variable objects before passing them to gradient functions.
Solution: Modified keras/src/backend/jax/core.py to automatically extract .value from Variable instances within the custom gradient decorator, eliminating the need for manual .value calls in user code. Added comprehensive tests in keras/src/ops/core_test.py::CoreOpsCorrectnessTest::test_custom_gradient_with_variable to prevent regression. This fix ensures seamless Variable handling across all backends without requiring workarounds.