[SPMD] Introduce high level manual sharding APIs (#6931)
Summary:
This pull request introduces:
1. enable_manual_sharding: which starts the manual sharding region.
2. disable_manual_sharding: which disable the manual sharding region.
Test Plan:
PJRT_DEVICE=TPU python test/spmd/test_xla_sharding.py -v -k test_manual_sharding_api_e2e