Repository: google/flax

Issue:
Error when using nn.scan with negative output_axes
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): N/A
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: ```flax==0.6.11, jax==0.4.9, jaxlib==0.4.9```
- Python version: ```3.8```
- GPU/TPU model and memory: N/A
- CUDA version (if applicable): N/A


### Problem you have encountered:

When using ```flax.linen.scan``` with a negative ```output_axes```, there is an unexpected ```AssertionError```. If I have understood the source code correctly, it is due to a typo [here](https://github.com/google/flax/blob/main/flax/core/axes_scan.py#L103) (namely, a minus sign instead of a plus sign). 

### What you expected to happen:

Apply scan as usual, stacking the outputs along the specified axis. 

### Logs, error messages, etc:

```
(projectabcde) lucaslingle@Lucass-MacBook-Pro projectabcde % python3 scripts/scan_issue.py
Traceback (most recent call last):
  File "scripts/scan_issue.py", line 39, in <module>
    main()
  File "scripts/scan_issue.py", line 32, in main
    params = cls().init(
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/linen/module.py", line 1689, in init
    _, v_out = self.init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/linen/module.py", line 1594, in init_with_output
    return init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/scope.py", line 968, in wrapper
    return apply(fn, mutable=mutable, flags=init_flags)({}, *args, rngs=rngs,
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/scope.py", line 936, in wrapper
    y = fn(root, *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/linen/module.py", line 2170, in scope_fn
    return fn(module.clone(parent=scope, _deep_clone=True), *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/linen/module.py", line 432, in wrapped_module_method
    return self._call_wrapped_method(fun, args, kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/linen/module.py", line 868, in _call_wrapped_method
    y = fun(self, *args, **kwargs)
  File "scripts/scan_issue.py", line 18, in __call__
    _, outputs = nn.scan(
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/linen/transforms.py", line 323, in wrapped_fn
    ret = trafo_fn(module_scopes, *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/lift.py", line 219, in wrapper
    y, out_variable_groups_xs_t = fn(
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/lift.py", line 806, in inner
    broadcast_vars, (carry_vars, c), (ys, scan_vars) = scanned(
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/axes_scan.py", line 151, in scan_fn
    ys = jax.tree_util.tree_map(transpose_from_front, out_axes, ys)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/jax/_src/tree_util.py", line 210, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/jax/_src/tree_util.py", line 210, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/axes_scan.py", line 106, in transpose_from_front
    return jax.tree_util.tree_map(trans, xs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/jax/_src/tree_util.py", line 210, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/jax/_src/tree_util.py", line 210, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/axes_scan.py", line 103, in trans
    assert pax < x.ndim
jax._src.traceback_util.UnfilteredStackTrace: AssertionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "scripts/scan_issue.py", line 39, in <module>
    main()
  File "scripts/scan_issue.py", line 32, in main
    params = cls().init(
  File "scripts/scan_issue.py", line 18, in __call__
    _, outputs = nn.scan(
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/axes_scan.py", line 151, in scan_fn
    ys = jax.tree_util.tree_map(transpose_from_front, out_axes, ys)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/axes_scan.py", line 106, in transpose_from_front
    return jax.tree_util.tree_map(trans, xs)
  File "/Users/lucaslingle/opt/miniconda3/envs/projectabcde/lib/python3.8/site-packages/flax/core/axes_scan.py", line 103, in trans
    assert pax < x.ndim
AssertionError
```



### Steps to reproduce:

```
# issue appears to be at https://github.com/google/flax/blob/main/flax/core/axes_scan.py#L101

import flax.linen as nn
import jax.random


class Foo(nn.Module):
    unused_config: int

    @nn.compact
    def __call__(self, state, input_dict):
        return state, nn.Dense(100)(input_dict["x"])


class Bar(nn.Module):
    @nn.compact
    def __call__(self, x):
        _, outputs = nn.scan(
            Foo,
            variable_broadcast="params",
            split_rngs=dict(
                params=False,
            ),
            in_axes=0,
            out_axes=-1,
        )(unused_config=123)(dict(unused_state_item=None), dict(x=x))
        return outputs


def main():
    cls = Bar
    params = cls().init(
        {"params": jax.random.PRNGKey(0)},
        jax.random.normal(jax.random.PRNGKey(1), shape=[8, 128, 16])
    )["params"]


if __name__ == "__main__":
    main()
```

Thank you for your attention to this matter!


Discussion:
I traced the assertion error `pax < x.ndim` in the debugger and it seems like `pax` has a value of 4, whereas `x.ndim` has a value of 3. `x.ndim` is 3 because your input data is 3-dimensional (with shape `[8, 128, 16]`). `pax` is derived [here](https://github.com/google/flax/blob/main/flax/core/axes_scan.py#L103), where `ax` is the `out_axes` input, which is -1 in this case. So `pax = x.ndim - ax = 3 - (-1) = 4`. Looking at this line of code, it seems like the assertion error will be raised if your `out_axes` argument is 0 or less. If you use `out_axes=2` instead, it should work. However I wonder if this is a bug in the code since if the `ax < 0` condition is fulfilled, it will be guaranteed that the assertion error will be raised.

cc: @marcvanzee @cgarciae 

Resolving patch:
diff --git a/flax/core/axes_scan.py b/flax/core/axes_scan.py
--- a/flax/core/axes_scan.py
+++ b/flax/core/axes_scan.py
@@ -100,7 +100,7 @@ def transpose_from_front(ax, xs):
 
     def trans(x):
       if ax < 0:
-        pax = x.ndim - ax
+        pax = x.ndim + ax
       else:
         pax = ax
       assert pax < x.ndim


Test patch: