Fix preprocess_function in run_summarization_flax.py
# What does this PR do?

`run_summarization_flax.py` has 

https://github.com/huggingface/transformers/blob/e7ed7ffdcb66c78d3437ed4c3a63c3640f50f436/examples/flax/summarization/run_summarization_flax.py#L535-L537

Using `jnp.array` here will cause `preprocess_function` to hang forever when it is used by `datasets.Dataset.map()` with `num_proc > 1`, when this script is running on a TPU VM.

I think it is related to #12719 and #12720

## Who can review?

@patil-suraj  @patrickvonplaten 

 

@patil-suraj - could you take this one? :-)