some issue in FlaxBertForMultipleChoice
Author: thevasudevgupta
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.7.0.dev0
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.1 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): Bert

## To reproduce

Steps to reproduce the behavior:

```python3
from transformers import BertConfig, FlaxBertForMultipleChoice
import numpy as np
model = FlaxBertForMultipleChoice(BertConfig())

model(np.ones((1, 4)))
```
<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

Output

```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/vasudevgupta/Local/transformers/src/transformers/models/bert/modeling_flax_bert.py", line 600, in __call__
    return self.module.apply(
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/flax/linen/module.py", line 936, in apply
    return apply(
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/flax/core/scope.py", line 687, in wrapper
    y = fn(root, *args, **kwargs)
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/flax/linen/module.py", line 1178, in scope_fn
    return fn(module.clone(parent=scope), *args, **kwargs)
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/flax/linen/module.py", line 275, in wrapped_module_method
    y = fun(self, *args, **kwargs)
  File "/Users/vasudevgupta/Local/transformers/src/transformers/models/bert/modeling_flax_bert.py", line 1022, in __call__
    reshaped_logits = logits.reshape(-1, num_choices)
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1322, in _reshape
    newshape = _compute_newshape(a, args[0] if len(args) == 1 else args)
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1316, in _compute_newshape
    return tuple(- core.divide_shape_sizes(np.shape(a), newshape)
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1316, in <genexpr>
    return tuple(- core.divide_shape_sizes(np.shape(a), newshape)
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/jax/core.py", line 1360, in divide_shape_sizes
    return handler.divide_shape_sizes(ds[:len(s1)], ds[len(s1):])
  File "/Users/vasudevgupta/miniconda3/lib/python3.8/site-packages/jax/core.py", line 1280, in divide_shape_sizes
    raise InconclusiveDimensionOperation(f"Cannot divide evenly the sizes of shapes {tuple(s1)} and {tuple(s2)}")
jax.core.InconclusiveDimensionOperation: Cannot divide evenly the sizes of shapes (1, 1) and (-1, 4)
```

I am probably missing something here. Please help me to figure out the issue here.

<!-- A clear and concise description of what you would expect to happen. -->
@LysandreJik @patrickvonplaten

Comments:
Ah I see! 

So multiple choice is actually the only class where the `input_ids` have to be of shape `(batch_size, num_choices, seq_length)` instead of `(batch_size, seq_length)` (see [docs](https://huggingface.co/transformers/model_doc/bert.html?highlight=flaxbert#flaxbertformultiplechoice) here).

This means that the code example should be changed to:

```python
from transformers import BertConfig, FlaxBertForMultipleChoice
import numpy as np
model = FlaxBertForMultipleChoice(BertConfig())

model(np.ones((1, 1, 4)))
```
in order to work :-)

Okay. I missed that😅. Sorry for this simple issue. I will fix bigbird also accordingly. 

Absolutely no worries ;-)