API to build non-sequential models

Hi there, and thanks for `flax`.

Flax's quite intuitive, but I noticed that it does not provide an API to handle branching in the model, allowing only sequential models.

I always find myself writing small patches, and I put together a quick PR as a proof of concept to test the idea. You can find it here: https://github.com/google/flax/pull/3243. Do you think it is something that you would consider adding in? Let me know if you like the idea, I can polish the PR.

Example usage below or in a [colab](https://colab.research.google.com/drive/1umHJIaK-AfyvTzfXkwmsAiC8X360W_2N?usp=sharing):
```python
module = nn.Sequential(
[
nn.Dense(4),
Split(2),
Parallel((nn.Dense(2), nn.Dense(2))),
Merge(jnp.sum)
]
)

```

Thanks!