GPT2 Flax "TypeError: JAX only supports number and bool dtypes, got dtype object in array"
On GPU

```
>>> from transformers import AutoTokenizer, FlaxAutoModelForCausalLM

>>> tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
>>> model = FlaxAutoModelForCausalLM.from_pretrained("gpt2-medium")
>>> input_context = "The dog"
>>> # encode input context
>>> input_ids = tokenizer(input_context, return_tensors="jax").input_ids
>>> # generate candidates using sampling
>>> outputs = model.generate(input_ids=input_ids, max_length=20, top_k=30, do_sample=True)

TypeError: JAX only supports number and bool dtypes, got dtype object in array
```

@patrickvonplaten @patil-suraj 
I see where this is coming from. For now, could you try initializing model and tokenizer like this
```
tokenizer = GPT2TokenizerFast.from_pretrained(model_id, padding_side="left", pad_token="<|endoftext|>")
model = FlaxGPT2LMHeadModel.from_pretrained(model_id, pad_token_id=50256,)
```

We'll soon publish a detailed colab about Flax generate <|||||>Done buddy. Worked. Thanks a lot. 
Does Flax ```model.generate``` makes use of caching ``` Query and Value ``` in attention layers ?

I have too run a benchmark for generation. Its fair to compare only if caching is supported. <|||||>yeah it only works with caching <|||||>Thanks . Closing the issue.