[Flax] Add scan_with_axes
Author: sanchit-gandhi
Pull request: yes
Labels: name: WIP

# What does this PR do?

Adds `scan_with_axes` to Flax Bert and its derived models.

TODO:
- [ ] Fix cookie cutter template
- [ ] Run `make fix-copies` (after review)

Fixes [#17399](https://github.com/huggingface/transformers/issues/17399)

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18341). All of your documentation changes will be reflected on that endpoint.

### Update

Current API:

1. With automatic init:
```python
model.scan_enable()  # to enable scan
model.scan_disable()  # to disable scan
```

2. Without automatic init:
```python
model.scan_enable()  # to enable scan in the nn.Module
params = model.convert_unroll_to_scan(params) # to convert the unrolled params to scan

model.scan_disable()  # to disable scan in the nn.Module (i.e. unrolled)
params = model.convert_scan_to_unroll(params) # to convert the scan params to unrolled
```

With automatic init, the params are converted from unrolled to scan under the hood ([`.convert_unroll_to_scan()`](https://github.com/huggingface/transformers/blob/98c42b01c2f1128b22bbda2cdc3d25e03b94af53/src/transformers/modeling_flax_utils.py#L435)). 

The params are converted to/from scan on the fly, with no memory overhead for conversion (_c.f._ [`.convert_unroll_to_scan()`](https://github.com/huggingface/transformers/blob/98c42b01c2f1128b22bbda2cdc3d25e03b94af53/src/transformers/modeling_flax_utils.py#L435)):
```python
for i in range(self.config.num_hidden_layers):
    # Stack the params for the N layers into one super block
    # and remove the unrolled layer params on the fly
    # -> no memory overhead for conversion!
    unrolled_layer = params.pop(key.replace("0", str(i)))
    stacked_params.append(unrolled_layer)
```

### Design question!

Should we include the unroll -> scan weight conversion in the `.from_pretrained()` method if `use_scan=True` is passed? Currently, setting `use_scan=True` will enable scan in the nn.Module, but will leave the params untouched. This is fine if loading scanned params from pre-trained. If loading unrolled params from pre-trained, the shape of the weights will not match those expected by the nn.Module (unrolled params, scanned nn.Module), meaning the weights will not be loaded!

Possible cases (loaded params, flag for `use_scan`):

| Params   | `use_scan` | Mismatch                           | Action                                    |
|----------|------------|------------------------------------|-------------------------------------------|
| Unrolled | False      | None                               | Fine!                                     |
| Unrolled | True       | params unrolled, nn.Module scanned | Requires conversion of params to scan     |
| Scan     | False      | params scan, nn.Module unrolled    | Requires conversion of params to unrolled |
| Scan     | True       | None                               | Fine!                                     |

Closing, see https://github.com/huggingface/transformers/issues/17399#issuecomment-5016868117