scan over a rolling window?

[PVarnai]

Hi!

I found some great examples of scanning over a sequence of data one by one, also using the lifted linen.scan version which is convenient with the flax training loop. But what if I want to scan a rolling window over a sequence of data? I couldn't really find anything, and I'm surprised this is not a usual operation (or it is I'm just missing something). For example, if I had a tensor x of shape (T, dims...), instead of getting the slices x[0, dims...], x[1, dims...], ... to operate on within each iteration of the scan, I would want to get x[0:window, dims...], x[1:window+1, dims...], and so on.

Thanks for any help!

[cgarciae]

Hey @PVarnai, we don't have a "sliding_window_scan", while it sounds cool you can mimic it by first creating the windows and then scanning over them. Here is a link to some JAX code for creating such windows:

jax-ml/jax#3171 (comment)

Copying the code here for convenience (credit to @erdmann):

from functools import partial
import jax
import jax.numpy as jnp
from jax import jit, vmap

@partial(jit, static_argnums=(1,))
def moving_window(a, size: int):
    starts = jnp.arange(len(a) - size + 1)
    return vmap(lambda start: jax.lax.dynamic_slice(a, (start,), (size,)))(starts)

a = jnp.arange(10)
print(moving_window(a, 4))

[[0 1 2 3]
 [1 2 3 4]
 [2 3 4 5]
 [3 4 5 6]
 [4 5 6 7]
 [5 6 7 8]
 [6 7 8 9]]

You can generalize it for ndarrays.

[PVarnai]

I see. Do I understand correctly that this creates copies of the window? So if I wanted to scan a window of length 100 for 200 time steps, this would create a memory of 200*100 numbers instead of just reusing the data from the 300 total time steps? Because that doesn't seem very feasible memory-wise, especially if you factor in some extra dimensions for the data and a batch size :/

Is it somehow possible to scan over a set of integers, passing all the data as a whole unscanned parameter, and creating the slices using the scanned integers for indexing the window?

[cgarciae]

As you state it might be very memory inefficient if the compiler does not optimize it (for some background have a read @mattjj response in the referenced issue), if you want to avoid the worst case scenario (probably recommended) you can instead scan over the window indices and slice the data inside the scan function accordingly e.g:

def f(xs, window_size):
  indices = jnp.arange(xs.shape[1] - window_size + 1)
  def body_fn(_, i):
    x = xs[:, i:i + window_size]
    y = ...
    return None, y
  _, ys = jax.lax.scan(body_fn, None, indices)
  return ys

[cgarciae]

You can adapt the previous idea to use it with Flax's nn.scan.