Add tree-mode jit (graph=False) to nnx.jit #5232
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
nnx.jit now accepts a `graph` parameter (default `True`). When `graph=False`, Modules are treated as regular JAX pytrees instead of going through the NNX graph protocol. This simpler mode behaves more like JAX as it assumes referential transparency (no sharing), it only propagates updates for Variables, roughly matching the expected behavior of mutable Hijax . Tree-mode removes the need for `graph.update_context` (the most complex part of NNX) and the NNX prefix/Lift APIs such as `StateAxes`, `StateSharding`, and `DiffState`. The implementation is thus much simpler, easier to maintain and optimize. Limitations of tree-mode (graph=False): - Shared Variable references are not supported (e.g. two sub-modules pointing to the same Linear layer will raise an error). - Mutable array references cannot be returned from the jitted function. - Capturing updates from the backward pass and forwarding state updates to captured objects is not available. The existing graph-mode (graph=True) remains the default and fully backward compatible. Users are not forced to migrate. PiperOrigin-RevId: 867790352
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add tree-mode jit (graph=False) to nnx.jit
nnx.jit now accepts a
graphparameter (defaultTrue). Whengraph=False,Modules are treated as regular JAX pytrees instead of going through the NNX
graph protocol. This simpler mode behaves more like JAX as it assumes referential transparency (no sharing), it only propagates updates for Variables, roughly matching the expected behavior of mutable Hijax .
Tree-mode removes the need for
graph.update_context(the most complexpart of NNX) and the NNX prefix/Lift APIs such as
StateAxes,StateSharding, andDiffState. The implementation is thus much simpler, easier to maintain and optimize.Limitations of tree-mode (graph=False):
pointing to the same Linear layer will raise an error).
to captured objects is not available.
The existing graph-mode (graph=True) remains the default and fully
backward compatible. Users are not forced to migrate.