Unroll a RNN core using output as the next input

[mmuravytskyi]

Hi guys, this is my first project in JAX and I'm in a bit of a struggle
I'm trying to unroll an RNN core having it to make new predictions based on its previous outputs. dynamic_unroll(core, ...) uses data from the passed input sequence and it's not what I want. The implementation below leads to trace leakages I'm not yet able to resolve. I'm wondering if there is a better way to make the unroll.

def sequence_loss(batch: Batch, rng_key: jnp.ndarray) -> jnp.ndarray:
    """Unrolls the network over a sequence of inputs & targets, gets loss."""
    core = make_network()
    batch_size, sequence_length, _ = batch['input'].shape  # (B, T, F)

    initial_state = core.initial_state(batch_size)

    _input = batch['input']  # (B, T, F)

    def body_fn(t: int, v :LoopValues) -> LoopValues:
        logit = v.logits[:, t, :]  # (B, F)
        new_logit, new_state = core(logit, v.state)  # (B, HIDDEN_SIZE)
        new_logit = jnp.reshape(new_logit, (BATCH_SIZE, 1, HIDDEN_SIZE))
        new_logit = hk.BatchApply(hk.Linear(FEATURE_NUM))(new_logit)
        # teacher forcing
        _key, _subkey = jax.random.split(v.rng_key)
        t_f_replace = jax.random.choice(_subkey, 
                jnp.array([1, 0], dtype=jnp.bool_), p = jnp.array([.1, .9]))
        t_f_keep = jnp.invert(t_f_replace)
        # (B, F)*() + (B, F)*()
        new_logit = jnp.reshape(new_logit, (BATCH_SIZE, FEATURE_NUM))
        t_f_logit = new_logit*t_f_replace + _input[:, t+1, :]*t_f_keep
        new_logits = v.logits.at[:, t + 1, :].set(t_f_logit)
        return LoopValues(logits=new_logits, state=new_state, rng_key=_key)

    output_seq = jnp.zeros((BATCH_SIZE, BLOCK_SIZE, FEATURE_NUM))
    output_seq = output_seq.at[:, 0, :].set(_input.at[:, 0, :].get())
    initial_values = LoopValues(logits=output_seq, state=initial_state, rng_key=rng_key)
    with jax.checking_leaks():
        lv: LoopValues = jax.lax.fori_loop(0, BLOCK_SIZE, body_fn, initial_values)

    loss = jnp.mean(jnp.abs(batch['target'] - lv.logits))  # MAE
    return loss

My implementation is based on https://github.com/deepmind/dm-haiku/blob/main/examples/rnn/train.py