[nnx] Arrays are state #3791
[nnx] Arrays are state #3791
Conversation
cgarciae
force-pushed
the
nnx-arrays-are-state
branch
from
f6a1846
to
f2e7de7
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
| str_path = '/'.join((*path, key)) | ||
| if value in ref_to_index: | ||
| variables.append((key, ref_to_index[value])) | ||
| if isinstance(value, Variable): |
There was a problem hiding this comment.
Nit: maybe move this into a separate branch before is_state_leaf to reduce the nesting here?
elif isinstance(value, Variable):
...
elif is_state_leaf(value):
...
| @@ -746,23 +756,32 @@ def _graph_update_dynamic( | |||
|
|
|||
| # case 2: subgraph is being updated | |||
| if is_node(current_value): | |||
| if isinstance(value, Variable): | |||
| if is_state_leaf(value): | |||
| raise ValueError( | |||
| f'Expected a subgraph for {key!r}, but got a Variable: {value!r}' | |||
There was a problem hiding this comment.
The error message needs updating.
| raise ValueError( | ||
| f'Expected a subgraph for {key!r}, but got a Variable: {value!r}' | ||
| ) | ||
| assert not isinstance(value, Variable) |
There was a problem hiding this comment.
Is this assert only needed to please the type checker?
flax/experimental/nnx/nnx/state.py
Outdated
|
|
||
| def __getattr__(self, key: Key) -> Variable | State: | ||
| def __getattr__(self, key: Key) -> tp.Any: |
There was a problem hiding this comment.
Is the return type genuinely Any here? It looks it could be StateLeaf | State where
StateLeaf = Union[Variable, jax.Array, np.ndarray]
There was a problem hiding this comment.
wow, coded this exact solution after implementing the TypeGuard :)
There was a problem hiding this comment.
...but hadn't seen you suggested exactly this
| @@ -152,6 +153,9 @@ def is_node(x: tp.Any) -> bool: | |||
| def is_node_type(x: type[tp.Any]) -> bool: | |||
| return x in NODE_TYPES or x is PytreeType | |||
|
|
|||
| def is_state_leaf(x: tp.Any) -> bool: | |||
There was a problem hiding this comment.
Can this be a typing.TypeGuard?
There was a problem hiding this comment.
great idea! I should start using TypeGuard more
| @@ -152,6 +153,9 @@ def is_node(x: tp.Any) -> bool: | |||
| def is_node_type(x: type[tp.Any]) -> bool: | |||
| return x in NODE_TYPES or x is PytreeType | |||
|
|
|||
| def is_state_leaf(x: tp.Any) -> bool: | |||
| return isinstance(x, (Variable, np.ndarray, jax.Array)) | |||
There was a problem hiding this comment.
Am I right that you're intentionally disallowing Python scalars here?
There was a problem hiding this comment.
Main problem here is that Variable is registered as a pytree so it would get accepted but we explicitly want to forbid it.
flax/experimental/nnx/nnx/state.py
Outdated
|
|
||
|
|
||
| def _state_unflatten( | ||
| static: tp.Tuple[Path, ...] | None, | ||
| static: tp.Tuple[Path, ...], |
There was a problem hiding this comment.
Use the built in tuple here and elsewhere.
cgarciae
force-pushed
the
nnx-arrays-are-state
branch
3 times, most recently
from
4ef0143
to
13c0c3a
Compare
cgarciae
force-pushed
the
nnx-arrays-are-state
branch
from
13c0c3a
to
8718110
Compare
What does this PR do?
jax.Arrayandnp.ndarrayare now treated as state same asVariables.Discussion
Previously we just treated all non-Variable and non-node as static, however, since we support pytrees and certain pytrees such as
opt_statein theOptimizerclass from #3782 sometimes include arrays such as the trainingsteparray, then we should also support Array. This sadly breaks from the simplicity ofStatebeing just a nested mapping withVariablesas leaves, but on the other hand Array is a very natural state type to try to support and in fact the NNX already supported it in the past.