[nnx] add transforms guide

[cgarciae]

What does this PR do?

Adds the Transforms guide.

Preview

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[IvyZX]

Thanks for making this guide! Super helpful and cool. Just a few nits on wordings.

[IvyZX]

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).

[IvyZX]

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.

[cgarciae]

Its a good point. I'll add a variation of the first example using vmap over __call__ so users know that its possible.

[IvyZX]

nit:

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.

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.

[IvyZX]

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)

[IvyZX]

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]

Thanks @IvyZX for the detailed feedback. I've integrated all the suggestions.