Repository: google/flax

Issue:
Implement a Sequential Module
Users often ask for this so it would be good to just add it. I can be as simple as this:

```python
class Sequential(nn.Module):
  layers: Sequence[nn.Module]

  def __call__(self, x):
    for layer in self.layers:
      x = layer(x)
    return x
```
Example usage:

```
class Foo(nn.Module):
  feature_sizes: List[int]
  @nn.compact
  def __call__(self, x):
    return Sequential([nn.Dense(sz, name=f'layers_{idx}')
                       for idx,sz in enumerate(self.feature_sizes)])(x)
```


Discussion:
Hi @marcvanzee,
I am interested in working on this issue. I submitted a PR #1156 if that is possible.

Resolving patch:
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -17,24 +17,26 @@
 
 # pylint: disable=g-multiple-import
 # re-export commonly used modules and functions
-from .activation import (celu, elu, gelu, glu, leaky_relu, log_sigmoid,
-                         log_softmax, relu, sigmoid, soft_sign, softmax,
-                         softplus, swish, silu, tanh, PReLU)
+from .activation import (PReLU, celu, elu, gelu, glu, leaky_relu, log_sigmoid,
+                         log_softmax, relu, sigmoid, silu, soft_sign, softmax,
+                         softplus, swish, tanh)
 from .attention import (MultiHeadDotProductAttention, SelfAttention,
-                        dot_product_attention, dot_product_attention_weights,
-                        make_attention_mask, make_causal_mask, combine_masks)
-from ..core import broadcast, DenyList, FrozenDict
+                        combine_masks, dot_product_attention,
+                        dot_product_attention_weights, make_attention_mask,
+                        make_causal_mask)
+from .combinators import Sequential
+from ..core import DenyList, FrozenDict, broadcast
+from .initializers import ones, zeros
 from .linear import Conv, ConvLocal, ConvTranspose, Dense, DenseGeneral, Embed
-from .module import (Module, compact, nowrap, enable_named_call,
-                     disable_named_call, override_named_call, Variable, init,
-                     init_with_output, apply, merge_param)
+from .module import (Module, Variable, apply, compact,
+                     disable_named_call, enable_named_call, init,
+                     init_with_output, merge_param, nowrap, override_named_call)
 from .normalization import BatchNorm, GroupNorm, LayerNorm
 from .pooling import avg_pool, max_pool, pool
-from .recurrent import GRUCell, LSTMCell, ConvLSTM, OptimizedLSTMCell
+from .recurrent import ConvLSTM, GRUCell, LSTMCell, OptimizedLSTMCell
 from .stochastic import Dropout
-from .transforms import (jit, named_call, checkpoint, remat, remat_scan, 
-                         scan, vmap, map_variables, vjp, jvp, custom_vjp,
+from .transforms import (checkpoint, custom_vjp, jit, jvp, map_variables,
+                         named_call, remat, remat_scan, scan, vjp, vmap,
                          while_loop)
-from .initializers import zeros, ones
 
 # pylint: enable=g-multiple-import
diff --git a/flax/linen/combinators.py b/flax/linen/combinators.py
new file mode 100644
--- /dev/null
+++ b/flax/linen/combinators.py
@@ -0,0 +1,39 @@
+"""Combinators of modules, such as a Sequential."""
+
+from typing import Callable, Sequence
+
+from flax.linen.module import Module
+
+class Sequential(Module):
+  """Applies a linear chain of Modules.
+
+  Meant to be used only for the simple case of fusing together callables where
+  the input of a particular module/op is the output of the previous one.
+
+  Modules will be applied in the order that they are passed in the constructor.
+
+  The apply() method of Sequential accepts any input and forwards it to the
+  first module it contains. It chains the output sequentially to the input of
+  the next module and returns the output of the final module.
+
+  Example usage::
+
+    class Foo(nn.Module):
+        feature_sizes: Sequence[int]
+
+        @nn.compact
+        def __call__(self, x):
+        return nn.Sequential([nn.Dense(layer_size, name=f'layers_{idx}')
+                              for idx, layer_size
+                              in enumerate(self.feature_sizes)])(x)
+  """
+  layers: Sequence[Callable]
+
+  def __call__(self, *args, **kwargs):
+    if not self.layers:
+      raise ValueError(f'Empty Sequential module {self.name}.')
+
+    outputs = self.layers[0](*args, **kwargs)
+    for layer in self.layers[1:]:
+      outputs = layer(outputs)
+    return outputs


Test patch:
diff --git a/tests/linen/linen_combinators_test.py b/tests/linen/linen_combinators_test.py
new file mode 100644
--- /dev/null
+++ b/tests/linen/linen_combinators_test.py
@@ -0,0 +1,93 @@
+"""Tests for flax.linen.combinators."""
+
+from typing import Any, Optional, Sequence
+
+from absl.testing import absltest
+
+from flax import linen as nn
+import jax
+from jax import numpy as jnp
+from jax import random
+import numpy as np
+
+# Parse absl flags test_srcdir and test_tmpdir.
+jax.config.parse_flags_with_absl()
+
+
+class MLP(nn.Module):
+  layer_sizes: Sequence[int]
+  activation: Optional[Any] = None
+  activation_final: Optional[Any] = None
+
+  @nn.compact
+  def __call__(self, inputs):
+    x = inputs
+    for layer_size in self.layer_sizes[:-1]:
+      x = nn.Dense(features=layer_size, kernel_init=nn.initializers.ones)(x)
+      if self.activation is not None:
+        x = self.activation(x)
+    x = nn.Dense(
+        features=self.layer_sizes[-1], kernel_init=nn.initializers.ones)(
+            x)
+    if self.activation_final is None:
+      return x
+    return self.activation_final(x)
+
+
+class SequentialTest(absltest.TestCase):
+
+  def test_construction(self):
+    sequential = nn.Sequential([nn.Dense(4), nn.Dense(2)])
+    key1, key2 = random.split(random.PRNGKey(0), 2)
+    x = random.uniform(key1, (3, 1, 5))
+    params = sequential.init(key2, x)
+    output = sequential.apply(params, x)
+    self.assertEqual(output.shape, (3, 1, 2))
+
+  def test_fails_if_layers_empty(self):
+    sequential = nn.Sequential([])
+    with self.assertRaisesRegex(ValueError,
+                                'Empty Sequential module'):
+      sequential.init(random.PRNGKey(42), jnp.ones((3, 5)))
+
+  def test_same_output_as_mlp(self):
+    sequential = nn.Sequential([
+        nn.Dense(4, kernel_init=nn.initializers.ones),
+        nn.Dense(8, kernel_init=nn.initializers.ones),
+        nn.Dense(2, kernel_init=nn.initializers.ones)
+    ])
+    mlp = MLP(layer_sizes=[4, 8, 2])
+
+    key1, key2 = random.split(random.PRNGKey(0), 2)
+    x = random.uniform(key1, (3, 5))
+    params_1 = sequential.init(key2, x)
+    params_2 = mlp.init(key2, x)
+
+    output_1 = sequential.apply(params_1, x)
+    output_2 = mlp.apply(params_2, x)
+    np.testing.assert_array_equal(output_1, output_2)
+
+  def test_same_output_as_mlp_with_activation(self):
+    sequential = nn.Sequential([
+        nn.Dense(4, kernel_init=nn.initializers.ones), nn.relu,
+        nn.Dense(8, kernel_init=nn.initializers.ones), nn.relu,
+        nn.Dense(2, kernel_init=nn.initializers.ones), nn.log_softmax
+    ])
+
+    mlp = MLP(
+        layer_sizes=[4, 8, 2],
+        activation=nn.relu,
+        activation_final=nn.log_softmax)
+
+    key1, key2 = random.split(random.PRNGKey(0), 2)
+    x = random.uniform(key1, (3, 5))
+    params_1 = sequential.init(key2, x)
+    params_2 = mlp.init(key2, x)
+
+    output_1 = sequential.apply(params_1, x)
+    output_2 = mlp.apply(params_2, x)
+    np.testing.assert_array_equal(output_1, output_2)
+
+
+if __name__ == '__main__':
+  absltest.main()