nnx.Sequential and nnx.split Problems

[DiagRisker]

Hello,

I recently tried less general architectures (compared to custom primitive ones), and I've been playing with nnx.Sequential to avoid writing a nnx.Module sub class for each variation. In the forward mode, everything works like expected, however I have some problem when extracting the model graph and parameter with nnx.split (in order to get a gradient)

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

x = jnp.zeros( (33,3))

@Partial(jax.jit, inline = True)
def conca(*args):
    return jnp.concatenate(args,axis=-1)
myNN = nnx.Sequential( nnx.Linear(3,12, rngs = nnx.Rngs(default = 0) ), Partial( conca, x) )
Y = myNN(x) ; print( Y.shape)

def loss( state, graphdef, rng_state, input):
    model = nnx.merge( graphdef, rng_state, state)
    return (model(input)**2).mean()
def Gjacob(f):
    """ General Gradient of respective parameters (projecting jacobian on canonical directions (tangent valued of ones)  
    """
    def Gacob(*x):
        y, vjp_fn = jax.vjp(f, *x)
        return vjp_fn(jnp.ones_like(y)) # tangent values at 1
    return Gacob


>>> graphdef, rng_state, state = nnx.split(myNN , nnx.RngState,...)
raise ValueError(
      f'Arrays leaves are not supported, at {path_str!r}: {value}'
    )
ValueError: Arrays leaves are not supported, at 'layers/1/0/0' 

Note that this problem does not occur with only Linear(nnx.Module) operators
I'm pretty sure someone in the community knows why it doesn't work :).
It would be a very handy option to have it working , since we could do less code for each module when doing different operations! (for instance here this is a way to do a residual bridge with concatenation)

Maybe some moderator know, @cgarciae?

Cheers,

[cgarciae]

Hey @DiagRisker, this should be resolved by #4612 soon. The issue is that currently NNX doesn't allow JAX Array as leaves, the solution would be to wrap x in Variable, e.g.

Partial(conca, nnx.Variable(x))

After #4612 gets merged you'll be able to use JAX Arrays directly.

[cgarciae]

Interesting, I'm not getting that error. Can you try updating JAX to the latest version? Here is the code I'm running:

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

x = jnp.zeros( (33,3))

@Partial(jax.jit, inline = True)
def conca(*args):
    return jnp.concatenate(args,axis=-1)
myNN = nnx.Sequential( nnx.Linear(3,12, rngs = nnx.Rngs(default = 0) ), Partial( conca, x) )
Y = myNN(x) ; print( Y.shape)

def loss( state, graphdef, rng_state, input):
    model = nnx.merge( graphdef, rng_state, state)
    return (model(input)**2).mean()

def Gjacob(f):
    """ General Gradient of respective parameters (projecting jacobian on canonical directions (tangent valued of ones)  
    """
    def Gacob(*x):
        y, vjp_fn = jax.vjp(f, *x)
        return vjp_fn(jnp.ones_like(y)) # tangent values at 1
    return Gacob

graphdef, rng_state, state = nnx.split(myNN , nnx.RngState,...)

gradsEB = Gjacob( Partial( loss, graphdef = graphdef, rng_state = rng_state, input = nnx.Variable(x)))(state)

print(jax.tree.map(jnp.shape, gradsEB))

Output:

(33, 15)
(State({
  'layers': {
    0: {
      'bias': VariableState(
        type=Param,
        value=(12,)
      ),
      'kernel': VariableState(
        type=Param,
        value=(3, 12)
      )
    },
    1: {
      0: {
        0: (33, 3)
      }
    }
  }
}),)

[DiagRisker]

I did python -m pip install --upgrade jax, flax (CPU install)

When executing your code with

import jax, jax.numpy as jnp, flax
from jax.tree_util import Partial
from flax import nnx
print(" JAX version : ", jax.__version__)
print(" flax version : ", flax.__version__)

JAX version : 0.5.3
flax version : 0.10.4
(33, 15)
Traceback (most recent call last):
File "", line 198, in _run_module_as_main
File "", line 88, in _run_code
File ... line ..., in
graphdef, rng_state, state = nnx.split(myNN , nnx.RngState,...)
~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
File "...\flax\nnx\graph.py", line 2042, in split
graphdef, flat_state = flatten(node)
~~~~~~~^^^^^^
File "...\flax\nnx\graph.py", line 502, in flatten
graphdef = _graph_flatten(
node,
...<6 lines>...
return_variables,
)
File "...\flax\nnx\graph.py", line 552, in _graph_flatten
nodedef = _graph_flatten(
value,
...<6 lines>...
return_variables,
)
File "...\flax\nnx\graph.py", line 552, in _graph_flatten
nodedef = _graph_flatten(
value,
...<6 lines>...
return_variables,
)
File "...\flax\nnx\graph.py", line 552, in _graph_flatten
nodedef = _graph_flatten(
value,
...<6 lines>...
return_variables,
)
File "...\flax\nnx\graph.py", line 591, in _graph_flatten
raise ValueError(
f'Arrays leaves are not supported, at {path_str!r}: {value}'
)
ValueError: Arrays leaves are not supported, at 'layers/1/0/0': [[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]]

Are you sure you pushed your latest branch?

[cgarciae]

can you also update jaxlib?

[cgarciae]

Ok, I checked again and you do get an error. However, #4612 is already merged! So you can now pass the JAX Arrays directly:

x = jnp.zeros( (33,3))

myNN = nnx.Sequential( nnx.Linear(3,12, rngs = nnx.Rngs(default = 0)), Partial(conca, x) )
Y = myNN(x) ; print( Y.shape)

graphdef, rng_state, state = nnx.split(myNN , nnx.RngState,...)

gradsEB = Gjacob( Partial( loss, graphdef = graphdef, rng_state = rng_state, input = x))(state)

I'll create a new release but you can install from main in the meantime via:

pip install git+https://github.com/google/flax.git

[DiagRisker]

Perfect, it works as you say! Many thanks, I am still interested in the official Pypi release, since the code should work in an enterprise environment (crucial for a PhD code)