How to use jax.lax.scan in Flax?

[marcvanzee]

(This question was asked by @marcosrdac on #16 , but I'm copying it here so it is easier accessible to users)

A Model instance currently fails when passed into a function like jit or scan. I have a fix for this but for now you could use scan(lambda c, x: lstm(c, x), c0, timeseries).

I plan to use Flax in my research on RNNs, but I'm struggling to understand some ideas behind Flax implementation for days. I really wanted to see an example of someone using the documented RNN cells, it would help me a lot!

Anyway, I tried to use the lambda trick, but it ended up not working, I get an error like this: ValueError: Jax transforms and modules cannot be mixed. when the cell is being defined. The definitions are like bellow (I don't know if that is that way it is mean't to be done, but I made it to be run on a single sample (i.e. x is a vector)):

class LRNNCell(nn.Module):
    @nn.compact
    def __call__(self, h, x):
        nh = h.shape[0]
        Whx = nn.Dense(nh)
        Whh = nn.Dense(nh, use_bias=False)
        Wyh = nn.Dense(1)

        h = nn.tanh(Whx(x) + Whh(h))
        y = nn.tanh(Wyh(h))
        return h, y

class LRNN(nn.Module):
    ny: Any
    nh: Any

    @nn.compact
    def __call__(self, x):
        h = jnp.zeros(self.nh)
        cell = LRNNCell()

        h, y = jax.lax.scan(lambda h, x: cell(h, x), h, x)
        return y[-self.ny:]

What am I missing?

[marcvanzee]

In Linen, we implemented "lifted" JAX transforms that allow you to apply JAX transformations to Modules directly, so you can define LRNNCell using that as follows:

class LRNNCell(nn.Module):
  @functools.partial(
      nn.transforms.scan,
      variable_broadcast='params',
      split_rngs={'params': False})
  @nn.compact
  def __call__(self, h, x):
      nh = h.shape[0]
      Whx = nn.Dense(nh)
      Whh = nn.Dense(nh, use_bias=False)
      Wyh = nn.Dense(1)

      h = nn.tanh(Whx(x) + Whh(h))
      y = nn.tanh(Wyh(h))
      return h, y

Both our sst2 and seq2seq are using the nn.transforms.scan transforms, so that may be helpful to get a better understanding as well.

[marcosrdac]

In Linen, we implemented "lifted" JAX transforms that allow you to apply JAX transformations to Modules directly, so you can define LRNNCell using that as follows:

class LRNNCell(nn.Module):
  @functools.partial(
      nn.transforms.scan,
      variable_broadcast='params',
      split_rngs={'params': False})
  @nn.compact
  def __call__(self, h, x):
      nh = h.shape[0]
      Whx = nn.Dense(nh)
      Whh = nn.Dense(nh, use_bias=False)
      Wyh = nn.Dense(1)

      h = nn.tanh(Whx(x) + Whh(h))
      y = nn.tanh(Wyh(h))
      return h, y

Both our sst2 and seq2seq are using the nn.transforms.scan transforms, so that may be helpful to get a better understanding as well.

I'm truly grateful for your help, @marcvanzee. I finally understand lifted transformations now! Thanks for pointing me the examples.

Now I have a code defining a linear RNN cell, meant to be run at a single sample (a vector).

class LRNNCell(nn.Module):
    @nn.compact
    def __call__(self, h, x):
        nh = h.shape[-1]
        Whx = nn.Dense(nh)
        Whh = nn.Dense(nh, use_bias=False)
        Wyh = nn.Dense(1)

        h = nn.tanh(Whx(x) + Whh(h))
        y = nn.tanh(Wyh(h))
        return h, y

Then I have a module to scan an input vector x containing various samples x[i] by running the RNN cell sequentially:

class LRNN(nn.Module):
    ny: Any
    nh: Any

    @nn.compact
    def __call__(self, x):
        cell = nn.scan(
            LRNNCell,
            variable_broadcast="params",
            split_rngs={"params": False},
        )

        h = jnp.zeros(self.nh)
        h, y = cell()(h, x[..., None])
        return y[-self.ny:, 0]

And a final lifted vmap on that class to transform our model from a single vector definition into one that can be run on datasets:

class BatchLRNN(nn.Module):
    ny: Any
    nh: Any

    @nn.compact
    def __call__(self, X):
        lrnn = nn.vmap(LRNN,
                       in_axes=0,
                       variable_axes={"params": None},
                       split_rngs={"params": False})
        return lrnn(ny=self.ny, nh=self.nh)(X)

I trained it on sinewave prediction and it works gracefully. Thanks, @marcvanzee !