Feature request: Add include_init_carry to RNN

[carlosgmartin]

RNN returns as the output sequence the sequence of new carries at each timestep, but excludes the initial carry. Inside RNN.__call__, the initial carry gets overwritten by the line carry, outputs = scan_output. Feature request: Add a flag include_initial_carry: bool that includes the initial carry in the output sequence.

[cgarciae]

Hey @carlosgmartin, RNN.__call__ only returns the last carry to the user. Are you getting the entire history of carry states? If so this is a bug. Can you share the code you are using?

[carlosgmartin]

@cgarciae Thanks for getting back to me. Here's an example of what I mean:

import jax
import flax

module = flax.linen.RNN(flax.linen.GRUCell(), 4)
key = jax.random.PRNGKey(0)
x = jax.random.uniform(key, (32, 20, 3))
w = module.init(key, x)
y = module.apply(w, x)
print(y.shape)  # (32, 20, 4)

A recurrent cell like GRU maps an input and state to a new state (the output). The output of the RNN is the sequence of new states. I was wondering if it would be possible to include the initial state in this sequence. The result would be a sequence of the same length as the input, plus 1. This is the state of the RNN at timestep zero, before it has seen any of the input sequence, which can be used to make predictions before any of the input has been seen.

[cgarciae]

You can get the initial carry by using the GRUCell.initialize_carry static method. Not sure we want to add this to the RNN API.

[carlosgmartin]

@cgarciae Like this? Would you suggest any changes vis-a-vis conforming to the flax API/style?

import jax
from jax import random, numpy as jnp
from flax import linen as nn


def prepend(item, seq, axis):
    return jax.tree_map(
        lambda item, seq: jnp.concatenate([jnp.expand_dims(item, axis), seq], axis),
        item,
        seq,
    )


class FullRNN(nn.RNN):
    def __call__(self, *args, **kwargs):
        y = super().__call__(*args, **kwargs)
        init_key = kwargs.get("init_key", None)
        if init_key is None:
            init_key = random.PRNGKey(0)
        y0 = self.cell.initialize_carry(init_key, y.shape[:-2], y.shape[-1])
        return prepend(y0, y, -2)


def main():
    module = FullRNN(nn.GRUCell(), 4)

    key = random.PRNGKey(0)

    key, subkey = random.split(key)
    x = random.uniform(subkey, (32, 20, 3))

    key, subkey = random.split(key)
    w = module.init(subkey, x)

    y = module.apply(w, x)
    print(y.shape)  # (32, 21, 4)


if __name__ == "__main__":
    main()

I'm using this in an autoregressive recurrent generative model where

• token N is stochastically generated from hidden state N
• hidden state N+1 is generated from hidden state N and token N

The initial (i.e., zeroth) token to be generated has no preceding tokens as input.

[cgarciae]

Not sure about using tree_map just to call jnp.concatenate, isn't y0 just an array?
Other than that it looks good.

[carlosgmartin]

Thanks for the feedback. I was considering cases like LSTMCell or ConvLSTMCell where it's not just an array.

[marcvanzee]

Converting this to a discussion since it seems the feature request is possible in the current API.