Fix pjit infeed test to work on CPU/TFRT and GPU.
Fix pjit infeed test to work on CPU/TFRT and GPU.

For these devices, the execution of the compiled function is
blocking, so transfer data to infeed before execution to avoid
deadlock.

google/jax,google/jax,google/jax,google/jax