Why isn't jax.lax.scan a transform?

[smorad]

Is there a reason why jax.lax.scan does not transform functions? Having jit, vmap operate on functions is super useful and very readable. It is a shame that it doesn't work for scan.

For example, ideally one could do something like:

def f(carry, x):
  return carry + x, carry + x

scan_f = jax.scan(f)
carry, output = scan_f(jnp.zeros(5,), jnp.arange(7 * 5).reshape(7, 5))

Or even use a decorator

@jax.scan
def f(carry, x):
  return carry + x, carry + x

carry, output = f(jnp.zeros(5,), jnp.arange(7 * 5).reshape(7, 5))

I suppose you can use partial, but it's less clean

def f(carry, x):
  return carry + x, carry + x

scan_f = partial(jax.lax.scan, f)
carry, output = scan_f(jnp.zeros(5,), jnp.arange(7 * 5).reshape(7, 5))

This is somewhat related to #23487

[smorad]

To add on, I think it would be ergonomic to include in_axes for scan just like vmap. Consider we often like to implement a recurrent function using a scan

theta = jnp.ones((3, 3))
x = jnp.ones((10, 3))
h = jnp.zeros((3,))

def f(carry, input):
  theta, x = input
  _, h = carry
  return h + theta @ x, None

scanf = jax.scan(f, in_axes=(0, None)) # sequence-map (smap?) over x but not params theta
scanf(init=h, xs=(theta, x))