Create Utils functions for PP

[Chenyaaang]

Description

The implementation for Pipeline Parallelism are splitted into the following small PRs.

• Util functions for PP on Jax (The current PR)
• worker changes
• runner changes
• platform changes
• torchax changes
• Jax changes
• Single host changes (this will be in vLLM)
• Multi host changes (Ray)

This PR is to create Util functions/classes for Jax's PP.

• jax_parallel_state.py is a simplified, Jax version of vLLM's parallel_state.py, to transfer intermediate tensors between each process, and also provide rank-control APIs for model loading, etc.
• jax_intermediate_tensor.py is a Jax version of vLLM's IntermediateTensor class. It contains the tensor dictionary that needs to be transfer from one PP rank to another.

Tests

E2E test has verified the whole PP implementation works properly.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.