[nnx] add cached_partial #4469
[nnx] add cached_partial #4469
Conversation
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
9 times, most recently
from
ad1a9a6 to
78be9e9
Compare
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
from
78be9e9 to
2ef1999
Compare
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
from
2ef1999 to
89aee13
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
15 times, most recently
from
186ca31 to
e421e3c
Compare
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
5 times, most recently
from
035857e to
8c13c94
Compare
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
from
8c13c94 to
dd5755a
Compare
| @@ -22,6 +22,7 @@ | |||
|
|
|||
|
|
|||
| class Config: | |||
| flax_use_flaxlib: bool | |||
There was a problem hiding this comment.
Do you still need this predefined attribute if you already have the bool_flag(...) line below?
If yes, does it make sense to add other flags here too?
There was a problem hiding this comment.
Python doesn't treat this as an attribute because Config is not a dataclass, this is a pure type hint.
flax/nnx/module.py
Outdated
| node_dict = graph.get_node_impl(self).node_dict(self) | ||
| node_impl = graph.get_node_impl(self) | ||
| if node_impl is None: | ||
| raise RuntimeError(f'Unsupported type: {type(self)}, this is a bug.') |
There was a problem hiding this comment.
Can we have better error message here?
There was a problem hiding this comment.
actually, changing this to an assert
flax/nnx/object.py
Outdated
| node_dict = graph.get_node_impl(node).node_dict(node) | ||
| node_impl = graph.get_node_impl(node) | ||
| if node_impl is None: | ||
| raise RuntimeError(f'Unsupported type: {type(node)}, this is a bug.') |
There was a problem hiding this comment.
And also better message here
There was a problem hiding this comment.
changed to an assert
flax/nnx/tracers.py
Outdated
| attributes={'jax_trace': self._jax_trace}, | ||
| path=path, | ||
| subtree_renderer=subtree_renderer, | ||
| import treescope # type: ignore[import-not-found,import-untyped] |
There was a problem hiding this comment.
Style nit: Can imports be moved to top level? We do top-level import of treescope in a bunch of files so adding this one won't make a difference.
| @@ -250,7 +250,9 @@ def test_nnx_to_linen(self): | |||
| assert y.shape == (1, 64) | |||
| np.testing.assert_allclose(y, x @ variables['params']['kernel']) | |||
| assert 'nnx' in variables | |||
| assert isinstance(variables['nnx']['graphdef'], nnx.GraphDef) | |||
| assert isinstance( | |||
| variables['nnx']['graphdef'], nnx.graph.NodeDef | nnx.graph.NodeRef | |||
There was a problem hiding this comment.
Should we have a type annotation for nnx.graph.NodeDef | nnx.graph.NodeRef?
There was a problem hiding this comment.
GraphDef is now defined as GraphDef = tp.Union[NodeDef[Node], NodeRef[Node]] but isinstance doesn't like it because its Generic.
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
5 times, most recently
from
95c54e5 to
7aea7d9
Compare
cgarciae
force-pushed
the
nnx-optimize-fingerprint
branch
from
7aea7d9 to
48c59d3
Compare
What does this PR do?
nnx.jit.flaxlib(unreleased)FLAX_USE_FLAXLIBflag.nnx.cached_partialAPI to cache the traversal or NNX objects, this API yields performance on par withsplit/merge.