[nnx] add transforms guide #4197
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
cgarciae
force-pushed
the
nnx-real-transforms-guide
branch
7 times, most recently
from
f0d5f35
to
f46c9ea
Compare
docs_nnx/guides/transforms.md
Outdated
| +++ | ||
|
|
||
| ### Graph updates propagation | ||
| JAX models inputs to transformations as trees, Flax NNX models inputs as graphs to allow for sharing references. However, to express most of Python's object model Flax NNX's state propagation machinery can track arbitrary updates to the objects as long as they're local (updates to globals inside transforms are not supported). This means that you can modify graph structure as needed, including updating existing attributes, adding/deleting attributes, swapping attributes, sharing (new) references between objects, sharing Variables between objects, etc. The sky is the limit! |
There was a problem hiding this comment.
JAX models inputs to transformations as trees, Flax NNX models inputs as graphs to allow for sharing references.
A bit hard to read - maybe:
JAX transformations see inputs as trees of arrays, and Flax NNX see inputs as graphs of Python references.
However, to express most of Python's object model Flax NNX's state propagation machinery can track arbitrary updates to the objects as long as they're local (updates to globals inside transforms are not supported).
This line also a bit verbose? Maybe just:
Flax NNX's state propagation machinery can track arbitrary updates to the objects as long as they're local to the input graph (updates to globals inside transforms are not supported).
| ) | ||
| x = jax.random.normal(random.key(1), (10, 2)) | ||
|
|
||
| def crazy_vector_dot(weights: Weights, x: jax.Array): |
There was a problem hiding this comment.
Not here, but I was hoping to see an example of transforming and using an nnx.Module method to showcase that it works and can be a natural pattern for users to take, since most transforms happen not at top level but in-between two layer definitions.
There was a problem hiding this comment.
Its a good point. I'll add a variation of the first example using vmap over __call__ so users know that its possible.
docs_nnx/guides/transforms.md
Outdated
| > With great power comes great responsibility. | ||
| > <br> \- Uncle Ben | ||
|
|
||
| While this feature is very powerful, it must be used with care as it can clash with JAX's underlying assumptions for certain transformations. For example, `jit` expects the structure of the inputs to be stable in order to cache the compiled function, changing the graph structure inside a `nnx.jit`-ed function cause continuous recompilations and performance degradation, `scan` on the other hand only allows a fixed `carry` structure, so adding/removing substates declared as carry will cause an error. |
There was a problem hiding this comment.
nit:
For example,
jitexpects the structure of the inputs to be stable in order to cache the compiled function, changing the graph structure inside annx.jit-ed function cause continuous recompilations and performance degradation,scanon the other hand only allows a fixedcarrystructure, so adding/removing substates declared as carry will cause an error.
For example, jit expects the structure of the inputs to be stable in order to cache the compiled function, so changing the graph structure inside a nnx.jit-ed function cause continuous recompilations and performance degradation. On the other hand, scan only allows a fixed carry structure, so adding/removing substates declared as carry will cause an error.
| ) | ||
| x = jax.random.normal(random.key(1), (10, 2)) | ||
|
|
||
|
|
There was a problem hiding this comment.
It's probably better to only call vmap once when only one call is needed, to avoid confusion. Same for the example below.
state_axes = nnx.StateAxes({nnx.Param: 0, Count: None}) # broadcast Count
@nnx.vmap(in_axes=(state_axes, 0), out_axes=1)
def stateful_vector_dot(weights: Weights, x: jax.Array):
assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'
assert x.ndim == 1, 'Batch dimensions not allowed'
weights.count += 1
return x @ weights.kernel + weights.bias
y = stateful_vector_dot(weights, x)
y = stateful_vector_dot(weights, x)
docs_nnx/guides/transforms.md
Outdated
| +++ | ||
|
|
||
| ### Random State | ||
| In Flax NNX random state is just regular state. This means that its stored inside Modules that need it and its treated as any other type of state. This is a simplification over Flax Linen where random state was handled by a separate mechanism. In practice Modules usually keep that need random state simply need a references to a `Rngs` object that is passed to them during initialization, and |
There was a problem hiding this comment.
In practice Modules usually keep that need random state simply need a references to a Rngs object that is passed to them during initialization, and use it to generate a unique key for each random operation.
What about:
In practice Modules simply need to keep a reference to a Rngs object that is passed to them during initialization, and use it to generate a unique key for each random operation.
cgarciae
force-pushed
the
nnx-real-transforms-guide
branch
from
f46c9ea
to
59acf38
Compare
cgarciae
force-pushed
the
nnx-real-transforms-guide
branch
from
59acf38
to
fb1a9cc
Compare
|
Thanks @IvyZX for the detailed feedback. I've integrated all the suggestions. |
What does this PR do?
Adds the
Transformsguide.Preview