fix(jax): improve JAX modules' names #5214
Conversation
Use `wraps` to keep the modules' names, so they won't be `FlaxModule`, which cannot be regonized. I realized it when implementing deepmodeling#5213.
Summary of ChangesHello @njzjz, 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 an issue where dynamically generated JAX/Flax modules were losing their original class names, making them difficult to identify and debug. By leveraging the Highlights
Changelog
Activity
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 issue of JAX modules not retaining their original names by correctly utilizing the functools.wraps decorator. The changes to type hints for the flax_module function are also a significant improvement, enhancing type correctness and clarity. The use of updated=() in the wraps decorator is a good practice when wrapping classes, preventing unintended overwrites of the wrapper's __dict__ while still copying essential metadata like __name__ and __module__ from the original class.
📝 WalkthroughWalkthroughType annotations for the Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~5 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
No actionable comments were generated in the recent review. 🎉 Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Updates the JAX backend’s flax_module decorator so wrapped nnx.Module classes retain the original module/class metadata, avoiding generic FlaxModule names that can interfere with recognition/introspection.
Changes:
- Add
functools.wraps(updated=())to preserve the wrapped class’s__name__/__qualname__/__module__/__doc__. - Refine
flax_moduletype annotations and docstring types to reflect class-in/class-out behavior.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| """Convert a NativeOP to a Flax module. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| module : NativeOP | ||
| module : type[NativeOP] | ||
| The NativeOP to convert. | ||
|
|
||
| Returns | ||
| ------- | ||
| flax.nnx.Module | ||
| type[flax.nnx.Module] | ||
| The Flax module. |
There was a problem hiding this comment.
The docstring still reads like module is an instance ("The NativeOP to convert"), but the function is used as a class decorator and now expects a class object. Updating the parameter/return descriptions to explicitly say "NativeOP subclass" / "Flax nnx.Module subclass" would avoid confusion for readers.
| @wraps(module, updated=()) | ||
| class FlaxModule(module, nnx.Module, metaclass=MixedMetaClass): | ||
| def __init_subclass__(cls, **kwargs: Any) -> None: |
There was a problem hiding this comment.
This change fixes a subtle introspection/serialization issue (module names no longer being FlaxModule). There doesn't appear to be a regression test that asserts @flax_module preserves __name__/__qualname__ (and avoids FlaxModule). Adding a small JAX unit test would help prevent future breakage.
| def flax_module( | ||
| module: NativeOP, | ||
| ) -> nnx.Module: | ||
| module: type[NativeOP], | ||
| ) -> type[nnx.Module]: |
There was a problem hiding this comment.
The return type annotation type[nnx.Module] loses the original class type (the returned class is also a subclass of the input module). This makes downstream typing inconsistent (e.g., decorated DPModel classes no longer type as their original DPModel base). Consider using a TypeVar so flax_module is typed as def flax_module(module: type[T]) -> type[T]: ... (and cast the generated class), optionally adding an nnx.Module Protocol/mixin if you want to preserve both facets.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5214 +/- ##
=======================================
Coverage 82.00% 82.00%
=======================================
Files 724 724
Lines 73801 73803 +2
Branches 3616 3615 -1
=======================================
+ Hits 60520 60522 +2
+ Misses 12120 12118 -2
- Partials 1161 1163 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
1 participant
Use
wrapsto keep the modules' names, so they won't beFlaxModule, which cannot be regonized. I realized it when implementing #5213.Summary by CodeRabbit