Rebase to tensorflow mainline

[merrymercy]

To avoid the divergence of our codebase and tensorflow/jax mainline, we want to rebase our repos to the tensorflow/jax mainline regularly. This brings several advantages:

1. Get the latest updates from the mainline. Google is actively improving Jax and XLA GPU backend.
2. Make future maintenance easier. If we do not sync with the mainline for a long time, we will never be able to sync with the mainline again due to too many conflicts.

I plan to rebase this repo once every two to three months following the tensorflow/jax releases. Note that rebases will rewrite the history of our repos, so we have to sync our local repos carefully after rebase.

Rebase instructions

https://docs.google.com/document/d/1_4X2Z4ewf2c9PcR2Ug4VpKNJxqbrtkhnHiQUt3yX2TI/edit?usp=sharing

Sync local repos

Ask @merrymercy for help

[merrymercy]

Rebase History

• Jul. 19, 2021
• Sept. 12, 2021
• Dec. 25, 2021 3fb790c
• Apr. 8, 2022 394a6cc
• Jul. 27, 2022 tf: b835ad5, jax:1a7c8831a8a37f76acc899bf6d6ef510349af044
• Nov. 20, 2022 tf: 736d0db, jax: 41417eed6f3770e05aeb1d2500967c5314d5de08

[GHGmc2]

@merrymercy jax had switched to use OpenXLA repo instead of TensorFlow: jax-ml/jax@172a831

[merrymercy]

@GHGmc2 Thanks for letting me know. We should follow this change. The maintenance team will discuss how to switch to OpenXLA.

[GHGmc2]

@merrymercy Thanks, I noticed that auto_sharding had been upstreamed to OpenXLA: https://github.com/openxla/xla/tree/main/xla/hlo/experimental/auto_sharding

Will the new code base reuse this part?