xla

[cherry-pick]Fix zero-1 bug for inferring local ranks (#5936)

#6237

Merged