Optionally Store Initial Scan Carry

[zombie-einstein]

jax.lax.scan is really handy for generating time-series from simulations or trajectories for RL training, but in these cases I often want to retain the initial carry state, which might look like

def step(state, _):
    new_state = ...
    return new_state, state

_, state_series = jax.lax.scan(step, initial_state, None, length=n_steps)

which produces a series n_steps long, the downside being you run the step one more time than you need to (which could be relatively expensive). There are ways around this that seem a bit less elegant, like storing the initial state and then joining it to the scan

_, state_series = jax.lax.scan(step, initial_state, None, length=n_steps-1)
state_series = jnp.concatenate([initial_state, state_series])

Would it be worth adding a flag to the scan to retain the initial state as the first element of its stacked output? Something like

def step(state, _):
    new_state = ...
    return new_state, new_state

_, state_series = jax.lax.scan(step, initial_state, None, length=n_steps-1, keep_init=True)

state_series here would have a length n_steps and start with the initial state, but we store the updated state at each step (and only run it as many times as we need).

[thisiscam]

The return value of the scanned function is a tuple of (carry, out), where carry and out do not have to be the same type. I think in your case, keep_init=True only makes sense if they are the same type.

So I would write a helper function just for this purpose, if it occurs very often:

def unfold(f, initial, length):
    def step(x, _):
      nx = f(x)
      return nx, nx
    _,  xs = jax.lax.scan(step, initial, None, length=length-1)
    return jnp.concatenate([initial, xs])

[zombie-einstein]

The return value of the scanned function is a tuple of (carry, out), where carry and out do not have to be the same type. I think in your case, keep_init=True only makes sense if they are the same type.

Ah yes a very good point about matching the carry type! Thanks for your suggestion also.