[Refactor] Move shared files to common

[kyuyeunk]

Description

Move files into common directory if they meet following requirements

• used in both jax and torchax path
• it is agnostic to either flax or torch and purely a jax functional code

Tests

https://buildkite.com/tpu-commons/tpu-inference-ci/builds/5315

some tests are failing but it's not relevant to this PR and will be fixed with this: #1091

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.

[jrplatin]

Thanks for this! Can you please make sure the docs are up-to-date with this refactor? I still see some old references

[kyuyeunk]

Thanks for this! Can you please make sure the docs are up-to-date with this refactor? I still see some old references

sgtm. given that the links there points to a files in main branch, I'll submit this commit first and create a PR that points to the newly modified files in main.