xla

45a9e0e9 - [SPMD] Allow mark_shard scalar tensors (#6101)

[SPMD] Allow mark_shard scalar tensors (#6101) Summary: Support mark_shard scalar tensors. Test Plan: python test/spmd/test_xla_sharding.py -v -k test_mark_shard_scalar