Flax BART training fails when evaluating
### System Info

- `transformers` version: 4.22.0.dev0
- Platform: Linux-5.15.0-41-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.9.1
- Flax version (CPU?/GPU?/TPU?): 0.6.0 (gpu)
- Jax version: 0.3.17
- JaxLib version: 0.3.15

### Who can help?

@sgugger @patil-suraj 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. make a dir `"./my_bart_model`
2. Train a tokenizer (let's use a small, Dutch corpus). **Note**: the repo README uses `tokenizer.save` but that only saves tokenizer.config and not the merges, so I think this is a second issue that should be fixed. Below I use `save_model` instead.

```python
from datasets import load_dataset
from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer

# load dataset
dataset = load_dataset("dbrd", "plain_text", split="train")

# Instantiate tokenizer
tokenizer = ByteLevelBPETokenizer()

def batch_iterator(batch_size=1000):
    for i in range(0, len(dataset), batch_size):
        yield dataset[i: i + batch_size]["text"]

# Customized training
tokenizer.train_from_iterator(batch_iterator(), vocab_size=50265, min_frequency=2, special_tokens=[
    "<s>",
    "<pad>",
    "</s>",
    "<unk>",
    "<mask>",
])

# Save files to disk
tokenizer.save_model("./my_bart_model")
```

3. Create a BART config for it

```python
from transformers import BartConfig
config = BartConfig.from_pretrained("facebook/bart-base", vocab_size=50265)
config.save_pretrained("./my_bart_model")
```

4. Train the model with a quick evaluation (command from the root of the transformers lib)

```sh
python examples/flax/language-modeling/run_bart_dlm_flax.py --output_dir ./my_bart_model --config_name ./my_bart_model --tokenizer_name ./my_bart_model --dataset_name dbrd --dataset_config_name plain_text --max_seq_length 128 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --learning_rate 1e-4 --warmup_steps 100 --overwrite_output_dir --logging_steps 200 --save_steps 500 --eval_steps 200
```

This leads to the following error (also note the VisibleDeprecation, although that might be unrelated to the triggered error):

```
transformers/examples/flax/language-modeling/run_bart_dlm_flax.py:288: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.
  {k: np.array([examples[i][k] for i in range(len(examples))]) for k, v in examples[0].items()}

Evaluating ...:  99%|██████████████████████████████████████████████████████████████████▏| 77/78 [00:08<00:00,  9.03it/s]Training...:  14%|█████████▍                                                         | 200/1429 [00:54<05:32,  3.69it/s]Epoch ... :   0%|                                                                                 | 0/3 [00:54<?, ?it/s]Traceback (most recent call last):
  File "transformers/examples/flax/language-modeling/run_bart_dlm_flax.py", line 964, in <module>
    main()
  File "transformers/examples/flax/language-modeling/run_bart_dlm_flax.py", line 896, in main
    model_inputs = data_collator(samples)
  File "transformers/examples/flax/language-modeling/run_bart_dlm_flax.py", line 291, in __call__
    batch["decoder_input_ids"] = shift_tokens_right(
  File "/home/bram/.local/share/virtualenvs/bart-tTDq1jwG/lib/python3.8/site-packages/transformers/models/bart/modeling_flax_bart.py", line 228, in shift_tokens_right
    shifted_input_ids[:, 1:] = input_ids[:, :-1]
IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed
```

### Expected behavior

No errors and preferably no deprecation warnings.
This seems to originate when this basic collate function is unsuccessful in creating valid numpy arrays, i.e., with one unspecified dimension, like `(2,)` instead of `(2, 16)` (bsz x seq_len).

https://github.com/huggingface/transformers/blob/65fb71bc762c46bb067306c1fd083b1cba87a095/examples/flax/language-modeling/run_bart_dlm_flax.py#L287-L289

I haven't dug further in why this occurs here, but presumably because the tokenizer is not explicitly padding to max length so in edge cases the sequences might not be of the same length? (The error mentioned above happened in the last batch (as you can see in the progress bar, 77/78)) so probably the last batch contained the last sequence which was not of the expected size.<|||||>Hey @BramVanroy! Thank you for posting this issue; the Flax BART training example is fresh off the press, so there might well be some small issues to fix.

Based on the traceback, I would agree that this is likely an issue related to tokenizer padding. Have you tried padding to max length? This seems like a very sensible next step!

Keep me posted with how you go on this! Happy to help if this remains a road-block.<|||||>I haven't figured this out yet. From reading the code, all blocks should be of size sequence length and the small remainder dropped:

https://github.com/huggingface/transformers/blob/cfd623a859890c6d106610d3c688064eadc7bd61/examples/flax/language-modeling/run_bart_dlm_flax.py#L657-L660

So I haven't looked further how this is being caused.

By the way, also getting these UserWarnings:

```
Some donated buffers were not usable: ShapedArray(float32[1,50265]), ShapedArray(float32[1026,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[1026,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[3072]), ShapedArray(float32[768,3072]), ShapedArray(float32[768]), ShapedArray(float32[3072,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768,768]), ShapedArray(float32[768]), ShapedArray(float32[768]), ShapedArray(float32[50265,768]).
See an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer-donation.
  warnings.warn(f"Some donated buffers were not usable: {', '.join(unused_donations)}.\n{msg}")
```

I am working on transposing `fairseq`'s data implementation to PyTorch and adding a full training example in transformers. Would you  be open to code-reviewing it when it's finished? The train script itself will probably borrow a lot from your code here if that's okay!<|||||>>  From reading the code, all blocks should be of size sequence length and the small remainder dropped:

Certainly, that should ideally be the case!

> By the way, also getting these UserWarnings:

Are those UserWarnings being thrown in the parameter update step? It suggests to me a mis-match between the update and parameter dtypes! 

> I am working on transposing fairseq's data implementation to PyTorch and adding a full training example in transformers. 

More than happy to perform a code-review when finished!

Also cc'ing @duongna21 who must take all credit for implementing the Flax BART training example!<|||||> Hi @BramVanroy, I'm the main author of `run_bart_dlm_flax`. It's unfortunate that I cannot reproduce your bug. Ran your code and things work as expected except for the `nan` loss bug that you can handle [like this](https://github.com/huggingface/transformers/pull/18458#discussion_r976303806).

![ảnh](https://user-images.githubusercontent.com/38061659/191475142-f568b80a-1dab-4358-96ff-686f719c8892.png)

After set `drop_last=True`, val loss is fine:

![ảnh](https://user-images.githubusercontent.com/38061659/191475435-e65ede77-5cf4-4247-8e3c-8b0f0abb9c80.png)

<|||||>@duongna21 Thanks for chiming in! I am going to close this as I am not working on this directly any more. I assume that drop_last should indeed fix the issue.