v0.11.0

[cgarciae]

v0.11.0 - Pytrees, MutableArrays, and more!

This version of Flax introduces some changes to improve interop with native JAX and adds support for the new jax.experimental.MutableArray. More on this soon! However, some breaking changes to align with the JAX way of doing things were necessary. Most code should remain intact, however, the following changes deviate from the current behavior:

• Rngs in standard layers: all standard layers no longer hold a shared reference to the rngs object given in the constructor, instead they now keep a fork-ed copy of the Rngs or RngStream objects. This impacts Using Rngs in NNX Transforms and Loading Checkpoints with RNGs.
• Optimizer Updates: the Optimizer abstraction no longer holds a reference to the model to avoid reference sharing, instead the model must be provided as the first argument to update.
• Modules as Pytrees: Modules are now pytrees! This avoid unnecessary use of split and merge when interacting trivially with raw JAX transforms (state must still be manually propagated if not using MutableArrays, and referential transparency is still an issue). This affects when operating on Pytrees containing NNX Objects with jax.tree.* APIs.

Checkout the full NNX 0.10 to NNX 0.11 migration guide.

In the near future we'll share more information about new ways of using NNX with JAX transforms directly by leveraging the new Pytree and MutableArray support. Stay tuned!

What's Changed

• [nnx] mutable array p3 by @cgarciae in [nnx] mutable array p3 #4755
• [nnx] allow method calls in ToLinen by @cgarciae in [nnx] allow method calls in ToLinen #4808
• Internal change by @copybara-service[bot] in Internal change #4807
• Preserve sharding information in axes_scan by @copybara-service[bot] in Preserve sharding information in axes_scan #4806
• Deduplicate contributing and philosophy and move to main site by @IvyZX in Deduplicate contributing and philosophy and move to main site #4809
• Fixed nnx.remat docstring rendering by @vfdev-5 in Fixed nnx.remat docstring rendering #4790
• Added a note to gemma guide about model's license consent on kaggle by @vfdev-5 in Added a note to gemma guide about model's license consent on kaggle #4776
• [nnx] ToLinen add abtract_init flag by @cgarciae in [nnx] ToLinen add abtract_init flag #4813
• Modify NNX to use id(variable) instead of nnx.Variables as dictionary by @divyashreepathihalli in Modify NNX to use id(variable) instead of nnx.Variables as dictionary #4814
• Allow using LazyRngs for flax init/apply. by @copybara-service[bot] in Allow using LazyRngs for flax init/apply. #4818
• [nnx] remove VariableState by @cgarciae in [nnx] remove VariableState #4800
• Fix failing CI jobs: trailing whitespace, deprecated .type usage by @vfdev-5 in Fix failing CI jobs: trailing whitespace, deprecated .type usage #4823
• [nnx] fix Rngs dtype check by @cgarciae in [nnx] fix Rngs dtype check #4820
• refactor: move usages of .value to [...] in modules_test.py by @lukeyeh in refactor: move usages of .value to [...] in modules_test.py #4815
• Added training script for Gemma model by @vfdev-5 in Added training script for Gemma model #4822
• [nnx] add flax_pytree_module flag by @cgarciae in [nnx] add flax_pytree_module flag #4811
• create ModelAndOptimizer symbol by @copybara-service[bot] in create ModelAndOptimizer symbol #4849
• [nnx] remove Optimizer.model attribute by @cgarciae in [nnx] remove Optimizer.model attribute #4842
• [nnx] add mutable array support in update by @cgarciae in [nnx] add mutable array support in update #4851
• Migrate transforms_test.py from .value to [...] by @lukeyeh in Migrate transforms_test.py from .value to [...] #4841
• 0.11.0 migration guide by @cgarciae in 0.11.0 migration guide #4854

New Contributors

• @divyashreepathihalli made their first contribution in Modify NNX to use id(variable) instead of nnx.Variables as dictionary #4814
• @lukeyeh made their first contribution in refactor: move usages of .value to [...] in modules_test.py #4815

Full Changelog: v0.10.7...v0.11.0

---

This discussion was created from the release v0.11.0.

[HeavyCrab]

Wait...
I remember the previous documentation (before 0.11.0) explicitly mentioned that Modules are not pytrees because shared references could be silently broken across JAX boundaries. Can you provide more explanation on this change? Like

1. Does making nnx.Modules pytrees now mean that shared references are no longer preserved (i.e., the previous concern still applies)?

2. Are split and merge now fully deprecated?

[acecchini]

I am really confused too. It seems now that nnx.Module is a pytree, however, from the doc they still use split and merge...

[HeavyCrab]

It seems that split and merge are preserved only for the circumstance where multiple nnx.Module objects with shared reference are involved. If x and y have shared reference, nnx.jit will deal with that, while jax.jit won't.

import jax
import jax.numpy as jnp
from flax import nnx

class mysubmodel(nnx.Module):
    def __init__(self):
        self.x = nnx.Variable(jnp.array(0))

class mymodel(nnx.Module):
    def __init__(self, submodel):
        self.submodel = submodel


submodel = mysubmodel()
model = mymodel(submodel)

def f(submodel, model):
    submodel.x += 1
    return submodel, model

jax_jit_f = jax.jit(f)
nnx_jit_f = nnx.jit(f)

print(nnx.state(jax_jit_f(submodel, model))) # the shared reference breaks when jax.jit is used
'''
State({
  0: {
    'x': Variable( # 1 (4 B)
      value=Array(1, dtype=int32, weak_type=True)
    )
  },
  1: {
    'submodel': {
      'x': Variable( # 1 (4 B)
        value=Array(0, dtype=int32, weak_type=True)
      )
    }
  }
})
'''

print(nnx.state((submodel, model))) # while this jax-jitted function is pure, so that the global vars won't be affected
'''
State({
  0: {
    'x': Variable( # 1 (4 B)
      value=Array(0, dtype=int32, weak_type=True)
    )
  }
})
'''

print(nnx.state(nnx_jit_f(submodel, model))) # the shared reference is preserved when nnx.jit is used
'''
State({
  0: {
    'x': Variable( # 1 (4 B)
      value=Array(1, dtype=int32, weak_type=True)
    )
  }
})
'''


print(nnx.state((submodel, model))) # the nnx-jitted function is not pure, so that the global vars will be affected
'''
State({
  0: {
    'x': Variable( # 1 (4 B)
      value=Array(1, dtype=int32, weak_type=True)
    )
  }
})
'''

But I'm not sure about this. Especially after the jax.array_ref API made it possible for a jax-jitted function to be impure. It's not clear that whether nnx.Module uses this, and whether we can say that when we pass a model into a JAX-jitted function, there will be no side effects.

[cgarciae]

@acecchini @HeavyCrab sorry for the confusion, I'm working on some guides around nnx.Pytree and ArrayRef but it has a been a bit slow due to other internal projects. The TLDR; is that for many use cases you will now be able to interface with JAX transforms directly as well as the jax.tree APIs, however pytree limitation still apply (no aliasing / reference sharing). split and merge sill serve very important use cases:

1. get explicit state groupings
2. preserve graph structure (in case sharing is needed).

Next priority is going to be documentation, to add new guides as well as update old guides to show how to use NNX directly with JAX transforms as well as some new interesting patterns.

[HeavyCrab]

In addition, I have another question. The release notes say that for some simple functions, we no longer need to use merge and split in the code. How is the efficiency in this case? Will there still be a relatively large overhead as mentioned in the earlier performance considerations? (I noticed that when we use jax.jit to decorate a function, the speed does seem to be noticeably faster than with nnx.jit.)

[vfdev-5]

Split and merge paradigm is still there and wont be deprecated. Here is an example of usage:

flax/examples/nnx_toy_examples/mutable_array_basic.py

Lines 62 to 71 in 2bf5748

	@jax.jit
	def train_step(model, optimizer, x, y):
	graphdef, params, counts = nnx.split(model, nnx.Param, Count)
	def loss_fn(params):
	model = nnx.merge(graphdef, params, counts)
	return jnp.mean((y - model(x)) ** 2)
	grads = jax.grad(loss_fn)(nnx.to_arrays(params))
	optimizer.update(model, grads)

In a special case when model contains only parameters, one can simplify the code and remove split/merge calls.
In general, with split and merge it is explicit what parameters are optimized.

[rademacher-p]

I'm sure the incoming docs will cover everything, but I'm curious about the Optimizer changes. Why avoid reference sharing? Sharing and updating the model was the big draw for me beyond pure optax