Fix preprocess_function in run_summarization_flax.py

[ydshieh]

What does this PR do?

run_summarization_flax.py has

transformers/examples/flax/summarization/run_summarization_flax.py

Lines 535 to 537 in e7ed7ff

	decoder_input_ids = shift_tokens_right_fn(
	jnp.array(labels["input_ids"]), config.pad_token_id, config.decoder_start_token_id
	)

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

[patrickvonplaten]

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

[patil-suraj]

Great catch! Thanks a lot for fixing this @ydshieh !