Repository: google/flax

Issue:
Initialization of Submodules Lifted with `flax.nn.scan`
One more issue 😄 . Promise this is the last one. There are a lot of question about `flax.nn.scan` and RTD and existing GitHub issues do not solve them.

With very deep model compilation times become insane and it takes about 1 hour to compile model for Nvidia runtime. So, I decided to prevent loop unrolling with `jax.lax.scan` and its lifting counterpart `flax.nn.scan`. However, I faced multiple issues. Incomplete list of issues follows.

1. There is no clear way to initialize scanned submodules. I came up with solution to pass everything as `args` and `kwargs` to `__call__` of submodule (in this case `MLP`).
2. There is no keyword argument of `flax.nn.scan` as RTD says.
3. Func `flax.nn.scan` always returns `(carry, args)` even if there is only `carry` and no `args`.
4. RTD says that `target` should be either a type of `nn.Module` or a function which accepts `nn.Module` (type?) as its first position argument.
5. If one specified `name` of modules in `MLP` then an exception is thrown. It is a bit strange because all parameter trees merged to a single parameter tree.

```python
import flax.linen as nn
import jax
import jax.numpy as jnp


def initializer(val):
    def init(key, shape, dtype):
        return jnp.full(shape, val, dtype)

    return init


class MLP(nn.Module):

    @nn.compact
    def __call__(self, xs, var):
        h = nn.Dense(features=2, kernel_init=initializer(var))(xs)
        h = nn.relu(h)
        h = nn.Dense(features=2, kernel_init=initializer(var))(xs)
        return xs + h, None


class Transformer(nn.Module):

    length: int = 3

    def setup(self):
        def fn(self, *args, **kwargs):
            return MLP(self, *args, **kwargs)

        # FAIL: Function instead of derived type from nn.Module does not work.
        #
        #   ScanMLP = nn.scan(target=fn, ...)
        #
        #   jax._src.traceback_util.UnfilteredStackTrace: TypeError:
        #   Transformer.setup.<locals>.fn() missing 1 required positional
        #   argument: 'self'

        # OK: No problems.
        ScanMLP = nn.scan(target=fn,
                          variable_axes={'params': 0},
                          variable_broadcast=False,
                          split_rngs={'params': True},
                          length=self.length)

        self.vars = jnp.arange(self.length)  # e.g. [0, 1, 2]
        self.mlp = ScanMLP()  # FAIL: ScanMLP(self.vars)

    @nn.compact  # OK: This decorator does nothing. Why?
    def __call__(self, xs):
        carry, out = self.mlp(xs, self.vars)  # OK: Axis 0 (implicitely).
        assert out is None
        return carry


model = Transformer(length=1250)
ys, state = jax.jit(model.init_with_output)(jax.random.PRNGKey(42),
                                            jnp.ones((3, 2)))
kernel = state['params']['mlp']['Dense_0']['kernel']
assert (kernel[0, ...] == jnp.zeros((2, 2))).all()
assert (kernel[1, ...] == jnp.ones((2, 2))).all()
```

In this experiments flax v0.6.3 and jax v0.4.1 are used.


Discussion:
Hey @daskol, sorry this took a while. I took you code and created a minimal working version:

```python
import flax.linen as nn
import jax
import jax.numpy as jnp

class MLP(nn.Module):
    @nn.compact
    def __call__(self, xs, _):
        h = nn.Dense(features=2)(xs)
        h = nn.relu(h)
        h = nn.Dense(features=2)(xs)
        return xs + h, None

class Transformer(nn.Module):
    n_layers: int = 4

    @nn.compact
    def __call__(self, x):
        ScanMLP = nn.scan(
            target=MLP, variable_axes={'params': 0}, variable_broadcast=False,
            split_rngs={'params': True}, length=self.n_layers)
        x, _ = ScanMLP()(x, None)
        return x

model = Transformer(n_layers=4)
y, variables = model.init_with_output(jax.random.PRNGKey(42), jnp.ones((1, 2)))

print(y.shape)
print(jax.tree_map(jnp.shape, variables))
```
What I think we can do is improve `nn.scan`'s documentation to show how to do this correctly, I am inclined to add a modified version of this example.
@cgarciae Thank you for your time. Eventually, I have managed to scan over modules with `@nn.compact` but another issue appears (see #2750 and underlying google/jax#13762). Briefly, the issue is that model initialization requires estimation some parameters (bias and variance of initializer) with non-jax code what can be done with `@nn.compact` decorator. Is it possible to rewrite you example in order to use `setup` instead of `@nn.compact`?
Hey @daskol, I am not sure I understand the need to use `setup` entirely (I did read the issue about `quad`) but in any case here is a `setup` version:

```python
import flax.linen as nn
import jax
import jax.numpy as jnp

class MLP(nn.Module):
    def setup(self):
        self.dense1 = nn.Dense(features=2)
        self.dense2 = nn.Dense(features=2)

    def __call__(self, xs, _):
        h = self.dense1(xs)
        h = nn.relu(h)
        h = self.dense2(h)
        return xs + h, None

class Transformer(nn.Module):
    n_layers: int = 4

    def setup(self):
        ScanMLP = nn.scan(
            target=MLP, variable_axes={'params': 0}, variable_broadcast=False,
            split_rngs={'params': True}, length=self.n_layers)
        self.scan_mlp = ScanMLP()
    
    def __call__(self, x):
        x, _ = self.scan_mlp(x, None)
        return x

model = Transformer(n_layers=4)
y, variables = model.init_with_output(jax.random.PRNGKey(42), jnp.ones((1, 2)))

print(y.shape)
print(jax.tree_map(jnp.shape, variables))
```
@cgarciae There are actually multiple issues with `flax.nn.scan`. I tried to solve the issue with my understanding of how `flax.nn.scan` works but faced to some unexpected behaviours and reported them here (FAIL and OK comment in the issue description).

However, the original problem which I reported  is that I need 

1. apply a `flax.nn.scan` to construct thousands of layers with special initialization parameters (variance);
2. and these parameters, in general, depend on the layer depth.

So, replacing `target=MLP` with `target=lambda: MLP(variance)` in your last snippet does not work and the code throws an exception. The full examples based on your code follows.

```python
import jax
import jax.numpy as jnp
import flax.linen as nn
from flax.linen.initializers import orthogonal


class MLP(nn.Module):

    var: float = 1.0

    def setup(self):
        self.dense1 = nn.Dense(features=2, kernel_init=orthogonal(self.var))
        self.dense2 = nn.Dense(features=2, kernel_init=orthogonal(self.var))

    def __call__(self, xs, _):
        h = self.dense1(xs)
        h = nn.relu(h)
        h = self.dense2(h)
        return xs + h, None

class Transformer(nn.Module):
    n_layers: int = 4

    def setup(self):
        # 1. This works.
        fn = MLP
        # 2. This does not work.
        def fn():
            return MLP(1 / self.n_layers)
        ScanMLP = nn.scan(target=fn, variable_axes={'params': 0},
            variable_broadcast=False, split_rngs={'params': True},
            length=self.n_layers)
        self.scan_mlp = ScanMLP()

    def __call__(self, x):
        x, _ = self.scan_mlp(x, None)
        return x

model = Transformer(n_layers=4)
y, variables = model.init_with_output(jax.random.PRNGKey(42), jnp.ones((1, 2)))

print(y.shape)
print(jax.tree_map(jnp.shape, variables))
```
UPD The most frequent exception thrown says that `self` is missing. I completely does not understand how it could be despite that I dug into `flax` sources and read a ton of code about how scopes work and how they are associated to modules and functional routines. It is still unclear for me why `self` is out of context captured by inner function `fn`.
@daskol There are 2 ways to do this, the **class** version: 

```python
class MLP(nn.Module):
    var: float = 1.0

    def setup(self):
        self.dense1 = nn.Dense(features=2, kernel_init=orthogonal(self.var))
        self.dense2 = nn.Dense(features=2, kernel_init=orthogonal(self.var))

    def __call__(self, xs, _):
        h = self.dense1(xs)
        h = nn.relu(h)
        h = self.dense2(h)
        return xs + h, None

class Transformer(nn.Module):
    n_layers: int = 4

    def setup(self):
        ScanMLP = nn.scan(target=MLP, variable_axes={'params': 0},
            variable_broadcast=False, split_rngs={'params': True},
            length=self.n_layers)
        self.scan_mlp = ScanMLP(1 / self.n_layers)

    def __call__(self, x):
        x, _ = self.scan_mlp(x, None)
        return x
```
And the **functional** version:
```python
class MLP(nn.Module):
    var: float = 1.0

    def setup(self):
        self.dense1 = nn.Dense(features=2, kernel_init=orthogonal(self.var))
        self.dense2 = nn.Dense(features=2, kernel_init=orthogonal(self.var))

    def __call__(self, xs):
        h = self.dense1(xs)
        h = nn.relu(h)
        h = self.dense2(h)
        return xs + h

class Transformer(nn.Module):
    n_layers: int = 4

    def setup(self):
        def scan_fn(mlp, x, _):
            return mlp(x), None
        self.scan = nn.scan(target=scan_fn, variable_axes={'params': 0},
            variable_broadcast=False, split_rngs={'params': True},
            length=self.n_layers)
        self.mlp = MLP(1 / self.n_layers)

    def __call__(self, x):
        x, _ = self.scan(self.mlp, x, None)
        return x
```
The general signature for either the Module's `__call__` method or the `scan_fn` function is 
```
(module, carry, xs) -> carry, ys
``` 
where in this case `xs = ys = None`. I will try to improve the example for the functional version in the docs for `nn.scan`.
@cgarciae Thank you very much! Now, it works perfectly.

> The general signature for either the Module's ...

I checked the ReadTheDocs. It says that the loop body should have the signature
```
(scope, body, carry, *xs) -> (carry, ys)
```
What is the difference? Does option `variable_broadcast` of `flax.nn.scan` affect the signature (assuming `scope` is the same as `body`)?

I found small typo in #2839 and left a comment. Could I ask you in advance to enable inter-sphinx links to `jax` and clarify what `vmap` or `scan` documentation says. For example, the current docs of `flax.linen.flax` say 

> To improve consistency with vmap, this version of scan uses in_axes and out_axes to determine which arguments are scanned over and along which axis.

It is unclear to me where `vmap`  refers to. To `jax.vmap` or to `flax.linen.vmap`?

Resolving patch:
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -738,47 +738,85 @@ def scan(target: Target,
 
   Example::
 
-    import flax
-    import flax.linen as nn
-    from jax import random
-
-    class SimpleScan(nn.Module):
-      @nn.compact
-      def __call__(self, c, xs):
-        LSTM = nn.scan(nn.LSTMCell,
-                       variable_broadcast="params",
-                       split_rngs={"params": False},
-                       in_axes=1,
-                       out_axes=1)
-        return LSTM()(c, xs)
-
-    seq_len, batch_size, in_feat, out_feat = 20, 16, 3, 5
-    key_1, key_2, key_3 = random.split(random.PRNGKey(0), 3)
-
-    xs = random.uniform(key_1, (batch_size, seq_len, in_feat))
-    init_carry = nn.LSTMCell.initialize_carry(key_2, (batch_size,), out_feat)
-
-    model = SimpleScan()
-    variables = model.init(key_3, init_carry, xs)
-    out_carry, out_val = model.apply(variables, init_carry, xs)
-
-    assert out_val.shape == (batch_size, seq_len, out_feat)
-
+    >>> import flax.linen as nn
+    >>> import jax
+    >>> import jax.numpy as jnp
+    ...
+    >>> class LSTM(nn.Module):
+    ...   features: int
+    ...
+    ...   @nn.compact
+    ...   def __call__(self, x):
+    ...     batch_size = x.shape[0]
+    ...     ScanLSTMCell = nn.scan(
+    ...       nn.LSTMCell, variable_broadcast="params",
+    ...       split_rngs={"params": False}, in_axes=1, out_axes=1)
+    ...
+    ...     carry = nn.LSTMCell.initialize_carry(
+    ...       jax.random.PRNGKey(0), (batch_size,), self.features)
+    ...     carry, x = ScanLSTMCell()(carry, x)
+    ...     return x
+    ...
+    >>> x = jnp.ones((4, 12, 7))
+    >>> module = LSTM(features=32)
+    >>> y, variables = module.init_with_output(jax.random.PRNGKey(0), x)
 
   Note that when providing a function to ``nn.scan``, the scanning happens over
   all arguments starting from the third argument, as specified by ``in_axes``.
-  So in the following example, the input that are being scanned over are ``xs``,
-  ``*args*``, and ``**kwargs``::
-
-    def body_fn(cls, carry, xs, *args, **kwargs):
-      extended_states = cls.some_fn(xs, carry, *args, **kwargs)
-      return extended_states
-
-    scan_fn = nn.scan(
-        body_fn,
-        in_axes=0,  # scan over axis 0 from third arg of body_fn onwards.
-        variable_axes=SCAN_VARIABLE_AXES,
-        split_rngs=SCAN_SPLIT_RNGS)
+  The previous example could also be written using the functional form as::
+
+    >>> class LSTM(nn.Module):
+    ...   features: int
+    ...
+    ...   @nn.compact
+    ...   def __call__(self, x):
+    ...     batch_size = x.shape[0]
+    ...
+    ...     cell = nn.LSTMCell()
+    ...     def body_fn(cell, carry, x):
+    ...       carry, y = cell(carry, x)
+    ...       return carry, y
+    ...     scan = nn.scan(
+    ...       body_fn, variable_broadcast="params",
+    ...       split_rngs={"params": False}, in_axes=1, out_axes=1)
+    ...
+    ...     carry = nn.LSTMCell.initialize_carry(
+    ...       jax.random.PRNGKey(0), (batch_size,), self.features)
+    ...     carry, x = scan(cell, carry, x)
+    ...     return x
+    ...
+    >>> module = LSTM(features=32)
+    >>> variables = module.init(jax.random.PRNGKey(0), jnp.ones((4, 12, 7)))
+
+  You can also use ``scan`` to reduce the compilation time of your JAX program
+  by merging multiple layers into a single scan loop, you can do this when
+  you have a sequence of identical layers that you want to apply iteratively
+  to an input. For example::
+
+    >>> class ResidualMLPBlock(nn.Module):
+    ...   @nn.compact
+    ...   def __call__(self, x, _):
+    ...     h = nn.Dense(features=2)(x)
+    ...     h = nn.relu(h)
+    ...     return x + h, None
+    ...
+    >>> class ResidualMLP(nn.Module):
+    ...   n_layers: int = 4
+    ...
+    ...   @nn.compact
+    ...   def __call__(self, x):
+    ...     ScanMLP = nn.scan(
+    ...       ResidualMLPBlock, variable_axes={'params': 0},
+    ...       variable_broadcast=False, split_rngs={'params': True},
+    ...       length=self.n_layers)
+    ...     x, _ = ScanMLP()(x, None)
+    ...     return x
+    ...
+    >>> model = ResidualMLP(n_layers=4)
+    >>> variables = model.init(jax.random.PRNGKey(42), jnp.ones((1, 2)))
+
+  To reduce both compilation and memory usage, you can use :func:`remat_scan`
+  which will in addition checkpoint each layer in the scan loop.
 
   Args:
     target: a ``Module`` or a function taking a ``Module``


Test patch: