xla

98df3d1a - [SPMD] Support mark_shard scalar tensors (#6158)

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