RNG count update during Manual Training

[DiagRisker]

Hello,

If you wanted to know how to update the RNG State count during manual Training. Here is an example (based on documentation MLP)

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

class MLP(nnx.Module): # Example Multi Layer Perceptron with dropout from documentation
    def __init__(self, rngs = nnx.Rngs(default = 54) ):
        self.linear = nnx.Linear(in_features = 3, out_features = 5, rngs = rngs)
        self.dropout = nnx.Dropout( .5, rngs = rngs)
    #
    def __call__(self,x):
        return self.dropout( self.linear(x) )

def loss(state, graphdef, rng_state, input): # arg to be differentiated with are 1st
    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

x = jnp.ones( ( 88, 3) )
model = MLP(); model.train()

# Let's verify the RNG state
graphdef, rng_state, state = nnx.split( model, nnx.RngState, ...) ; print( rng_state) 

State({
'dropout': {
'rngs': {
'default': {
'count': VariableState( # 1 (4 B)
type=RngCount,
value=Array(2, dtype=uint32),
tag='default'
),
'key': VariableState( # 1 (8 B)
type=RngKey,
value=Array((), dtype=key) overlaying:
[ 0 54],
tag='default'
)
}
}
}
})

# When differentiating, the model is not called and updated in RNG state (even if vjp returns Jacobian and forward output)
gradsEB = Gjacob( Partial( loss, graphdef = graphdef, rng_state = rng_state))(state) # --> no RNG count state
graphdef, rng_state, state = nnx.split( model, nnx.RngState, ...) ; print( rng_state) 

State({
'dropout': {
'rngs': {
'default': {
'count': VariableState( # 1 (4 B)
type=RngCount,
value=Array(2, dtype=uint32),
tag='default'
),
'key': VariableState( # 1 (8 B)
type=RngKey,
value=Array((), dtype=key) overlaying:
[ 0 54],
tag='default'
)
}
}
}
})

grads = nnx.value_and_grad(lambda *arg,**kar: loss(*arg , graphdef = graphdef, rng_state = rng_state, input = x,**kar))( state)
graphdef, rng_state, state = nnx.split( model, nnx.RngState, ...) ; print( rng_state) 

State({
'dropout': {
'rngs': {
'default': {
'count': VariableState( # 1 (4 B)
type=RngCount,
value=Array(2, dtype=uint32),
tag='default'
),
'key': VariableState( # 1 (8 B)
type=RngKey,
value=Array((), dtype=key) overlaying:
[ 0 54],
tag='default'
)
}
}
}
})

Y = model(x) # -> updates RNG count state
graphdef, rng_state, state = nnx.split( model, nnx.RngState, ...) ; print( rng_state) 

State({
'dropout': {
'rngs': {
'default': {
'count': VariableState( # 1 (4 B)
type=RngCount,
value=Array(3, dtype=uint32),
tag='default'
),
'key': VariableState( # 1 (8 B)
type=RngKey,
value=Array((), dtype=key) overlaying:
[ 0 54],
tag='default'
)
}
}
}
})

We can see that no count change occured for the gradient pass. This is troublesome for mini batch processing if dropout does not update its rng count during each batch. How to force it to update? -> jax.tree.map is an option but it does not directly work with nnx Filters!
Based on https://flax.readthedocs.io/en/latest/guides/randomness.html, you can write a shortcut :

# extract rng count solely
RngC = nnx.state(model, nnx.RngCount) # only counts
print( RngC)

def UpdateRNG(model):
    graphdef, RNGCount, state = nnx.split(model , nnx.RngCount,...)  # RNGCount = tree_map(lambda x : x+1, RNGCount)
    return nnx.merge( graphdef, tree_map(lambda x : x+1, RNGCount), state )
    # nnx.update(model, tree_map(lambda x : x+1, RNGCount) ) # this is also an option

model = UpdateRNG(model)
RngC = nnx.state(model, nnx.RngCount) # only counts
print( RngC)

State({
'dropout': {
'rngs': {
'default': {
'count': VariableState( # 1 (4 B)
type=RngCount,
value=Array(4, dtype=uint32),
tag='default'
),
'key': VariableState( # 1 (8 B)
type=RngKey,
value=Array((), dtype=key) overlaying:
[ 0 54],
tag='default'
)
}
}
}
})

This is a simple and useful way to control how rng goes during the training
If the rng state have different update methods (than count +=1 ), this will not solve the problem.
In this case, I leave contributors the option to complete this case