Dense + while loop interaction

[gahdritz]

This is very similar to #2920, where I was trying the same thing but with Dropout.

The following toy example, where I'm trying to run a Dense layer several times (with weight sharing) in a while loop, fails:

import jax
import jax.numpy as jnp
from flax import linen as nn

class MyModel(nn.Module):
  num_neurons: int

  @nn.compact
  def __call__(self, x, training: bool):
    dense = nn.Dense(self.num_neurons, name="dense")
    
    def fn(mdl, x):
        x = dense(x)
        return x

    body = lambda mdl, inp: (
        inp[0] + 1,
        jax.lax.stop_gradient(fn(mdl, inp[1])),
    )

    _, x = nn.while_loop(
        lambda _, x: x[0] < 3,
        body,
        self,
        (0, jnp.zeros([*x.shape])),
    )

    x = dense(x)

    return x

root_key = jax.random.PRNGKey(seed=0)
main_key, params_key = jax.random.split(key=root_key, num=2)

my_model = MyModel(num_neurons=3)
x = jnp.empty((3, 4, 4))
variables = my_model.init(params_key, x, training=False)
params = variables['params']

y = my_model.apply({'params': params}, x, training=True)

It crashes with the following error:

Traceback (most recent call last):
  File "/n/home12/gahdritz/projects/proj/test.py", line 37, in <module>
    variables = my_model.init(params_key, x, training=False)
  File "/n/home12/gahdritz/projects/proj/test.py", line 21, in __call__
    _, x = nn.while_loop(
  File "/n/home12/gahdritz/projects/proj/test.py", line 18, in <lambda>
    jax.lax.stop_gradient(fn(mdl, inp[1])),
  File "/n/home12/gahdritz/projects/proj/test.py", line 13, in fn
    x = dense(x)
  File "/n/home12/gahdritz/projects/proj/venv/lib/python3.9/site-packages/flax/linen/linear.py", line 196, in __call__
    kernel = self.param('kernel',
  File "/n/home12/gahdritz/projects/proj/venv/lib/python3.9/site-packages/flax/core/tracers.py", line 36, in check_trace_level
    raise errors.JaxTransformError()
flax.errors.JaxTransformError: Jax transforms and Flax models cannot be mixed. (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.JaxTransformError)

Any tips?

[cgarciae]

Hey @gahdritz, the error is a bit uninformative. Same as with nn.scan you have to pass broadcast_variables='params' to indicate this collection will be shared, and you should also pass split_rng for initialization. Here is the working example:

import jax
import jax.numpy as jnp
from flax import linen as nn

class MyModel(nn.Module):
  num_neurons: int

  @nn.compact
  def __call__(self, x, training: bool):

    dense = nn.Dense(self.num_neurons, name="dense")

    if self.is_initializing():
        dense(x)
    
    def fn(mdl, x):
        x = dense(x)
        return x

    body = lambda mdl, inp: (
        inp[0] + 1,
        jax.lax.stop_gradient(fn(mdl, inp[1])),
    )

    _, x = nn.while_loop(
        lambda _, x: x[0] < 3,
        body,
        self,
        (0, x),
        broadcast_variables='params',
        split_rngs={'params': False},
    )

    return x

root_key = jax.random.PRNGKey(seed=0)
main_key, params_key = jax.random.split(key=root_key, num=2)

my_model = MyModel(num_neurons=4)
x = jax.random.normal(key=main_key, shape=(3, 4, 4))
variables = my_model.init(params_key, x, training=False)
params = variables['params']

y = my_model.apply({'params': params}, x, training=True)
print(y)

One thing that surprised me is that I had to condition on self.is_initializing to run the layer before running the loop. nn.scan automatically initializes broadcasts on the first pass, seems this transformation is missing that feature. I'll create an issue.

[gahdritz]

Thanks again!