Improve lifted transforms docstrings

(After feedback from @zhangqiaorjc)

Some `nn.scan` arguments are under-documented / make undocumented assumptions. In the example below, the scanning actually happens over the third argument of `body_fn` (`sub` and `layer_in` are ignored):

```py
def body_fn(sub, layer_in, layer_states):
# some code

scan_fn = nn.scan(
body_fn,
in_axes=0, # scan over axis 0 for layer_states only
variable_axes=SCAN_VARIABLE_AXES,
split_rngs=SCAN_SPLIT_RNGS)
```

However, this is not made clear in the [current documentation](https://flax.readthedocs.io/en/latest/_autosummary/flax.linen.scan.html):

**in_axes**: Specifies the axis to scan over for the arguments. Should be a prefix tree of the arguments. Use `flax.core.broadcast` to feed an entire input to each iteration of the scan body.

We should go over all arguments of `nn.scan.` and check whethey they make any assumptions that should be mentioned.