Support pytree classes with dunder-style flatten/unflatten methods

[jakevdp]

Fixes #33199

[gemini-code-assist]

Summary of Changes

Hello @jakevdp, I'm Gemini Code Assist¹! 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 enhances JAX's pytree utility by allowing users to define flatten and unflatten methods using dunder (double underscore) prefixes, such as __tree_flatten__ and __tree_unflatten__. This change provides greater flexibility in naming conventions for custom pytree nodes, addressing potential conflicts and offering a more Pythonic way to implement these protocols. The update ensures backward compatibility while expanding the options for pytree registration.

Highlights

• Dunder Method Support: Introduced support for dunder-style __tree_flatten__ and __tree_unflatten__ methods within register_pytree_node_class for custom pytree nodes.
• Keyed Pytree Dunder Support: Extended dunder method support to register_pytree_with_keys_class, allowing __tree_flatten_with_keys__ and __tree_unflatten__.
• Documentation Update: Updated docstring examples for register_pytree_node_class to reflect the new dunder method naming convention.
• New Test Cases: Added new test classes (SpecialDunder and SpecialDunderWithKeys) and integrated them into existing test suites to validate the correct behavior of pytree classes using the new dunder methods.

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

1. 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. ↩

[gemini-code-assist]

Code Review

This pull request adds support for dunder-style __tree_flatten__ and __tree_unflatten__ methods for pytree registration, which is a nice pythonic addition. The implementation looks correct, and the tests cover the new functionality.

I have a few suggestions to make the API more robust by enforcing consistency when using dunder methods. Currently, it's possible to mix dunder and non-dunder methods (e.g., __tree_flatten__ with tree_unflatten), which could be confusing. My suggestions would make the API stricter and less error-prone. I also found a minor style issue with trailing whitespace.