lax.scan with generator

[pharringtonp19]

I was naively hoping that we could pass a generator as the state to lax.scan. The aim was to use lax.scan to perform "mini-batch training". I would be curious to know of any workarounds/alternatives. -- Thanks

A minimal example which captures the error is reproduced below:

import jax 
import jax.numpy as jnp 

X = jax.random.normal(jax.random.PRNGKey(0), (10, 2))
batch_size = 2

def sampler(X, batch_size=2):
  for i in range(len(X)//batch_size):
    yield X[i*batch_size: (i+1)*batch_size, :]

def inner_update(carry, t):
  sampler = carry 
  return sampler, jnp.max(next(sampler))

jax.lax.scan(inner_update, sampler(X, batch_size), jnp.arange(len(X)//batch_size)) 

TypeError: Value <generator object sampler at 0x7f0d1826ccd0> with type <class 'generator'> is not a valid JAX type

[jakevdp]

No, lax.scan does not accept generator expressions, mainly because JAX code is lowered to XLA and XLA does not have any concept of generator expressions. You can accomplish roughly the same thing by putting the generator expression logic in the update function:

import jax
import jax.numpy as jnp 

X = jax.random.normal(jax.random.PRNGKey(0), (10, 2))
batch_size = 2

def inner_update(carry, t):
  i, X = carry
  batch = jax.lax.dynamic_slice(X, (i * batch_size, 0), (batch_size, X.shape[1]))
  return (i + 1, X), jnp.max(batch)

_, out =jax.lax.scan(inner_update, (0, X), jnp.arange(len(X)//batch_size))

[cgarciae]

Hey @jakevdp ! I've been thinking about this problem, in case that the generator is a TF Dataset or PT Dataloader what would be the recommended solution?

You could accumulate N batches into a "super batch" and use scan on that, but it might not payoff, should depend a lot on how fast you can generate these super batches.

[shoyer]

@shoyer Would this second suggestion (using jax.experimental.host_callback.call) allow one to compile the entire training loop?

Yes

[cgarciae]

An alternative would be to use jax.experimental.host_callback.call to pass data inside an ongoing XLA computation

Very interesting, but is there an actual speedup given the Python interpreter is involved?

[shoyer]

Very interesting, but is there an actual speedup given the Python interpreter is involved?

If you're using TF datasets and compiling the entire training loop, then very little time should be spent back in the Python interpreter. The Python GIL is not really a problem. (This is true for both the for loop and host_callback options.)

[cgarciae]

@shoyer On CPU + a very quick generator, scan + call was ~ 5x slower that a python for loop. Tried it like this (here model is a pytree):

def train_step(model: Model, _) -> tp.Tuple[Model, Logs]:
    x, y = hcb.call(
        train_sampler,
        (),
        result_shape=result_shape,
    )
    model, train_logs = model.train_step(x, y)

    return model, train_logs

model, train_logs = jax.lax.scan(train_step, model, None, length=n_batches)

[shoyer]

I'm not too surprised to see that Python control flow is faster, that path is certainly the one that has been most optimized.