fix the issue when eager mode jax is triggered in undesired places

[Hanjun-Dai]

Existing implementation will trigger the jax eager mode, which compiles many small pieces of code on-demand (in order to get it run on TPU) at every training iteration. This will cause several issues:

Before the fix:

• We see a lot of complains from wandb: WARNING Tried to log to step 0 that is less than the current step xxx. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
  This is because the jax compilation message also gets monitored and logged into wandb. But jax doesn't know the step and default step 0 is used. That's why it prints a lot of these warnings during training.

• TPU utilization: the eager mode especially the re-compilation at every iteration takes time, which can make things slow and lower the TPU utilization

• Probably too many cached binaries will be created until all kinds of shapes have been compiled.

After the fix:

• wandb warning is gone

• TPU utilization is higher and it should run faster

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[wang2yn84]

There are couple of other callersites of this utility function. E.g. the vllm_sampler and sglang_jax_sampler wraps prompt id as Jax array. Can you update them accordingly. The other callersites such as dpo_trainer you can leave them there. We need to refactor them to use the generate/utils.py.

[wang2yn84]

Actually let me just merge this PR, it's very important to us and I'd like to land it ASAP. I will fix the rest from internal.