What does linen.transforms.scan do?

[cccntu]

I saw it here:

flax/examples/seq2seq/train.py

Lines 157 to 164 in 93078a4

	class EncoderLSTM(nn.Module):
	@functools.partial(
	nn.transforms.scan,
	variable_broadcast='params',
	split_rngs={'params': False})
	@nn.compact
	def __call__(self, carry, x):
	return nn.LSTMCell()(carry, x)

And it's defined here with a TODO:

flax/flax/linen/transforms.py

Lines 198 to 202 in 93078a4

	# TODO: provide wrappers with annotated args/kwargs and docstrings.
	vmap = functools.partial(lift_transform, lift.vmap)
	jit = functools.partial(lift_transform, lift.jit)
	remat = functools.partial(lift_transform, lift.remat)
	scan = functools.partial(lift_transform, lift.scan)

[avital]

I believe that @jheek was planning on adding clear docstrings to these.

[jheek]

These are the "lifted" jax transformations for linen. The short story is that jax transforms don't work natively with stateful things (eg. variables, rng sequences) like linen Modules. These lifted transforms wrap the existing jax transforms and add functionality on top it to support modules fully. You can check flax.core.lift for their implementation and a short docstring. Over the last months we have stabilised the implementation and you can soon expect proper documentation and examples for these features.