[SPMD] Support manual sharding (#6915)
Summary:
This pull request makes SPMD support the manual sharding type via a new private API called: _mark_manual_sharding. I don't expect users will need to call this function explicitly.
Besides adding support for the sharding annotation, we also need to define the behavior of the data shards. For data, the current behavior is error out.
Test Plan:
PJRT_DEVICE=TPU python test/spmd/test_xla_sharding.py -v -k test__mark_manual_sharding
