Use source tensor sizes after the sharding created so far when tie-breaking between candidates during explicit reshards.

[copybara-service]

Use source tensor sizes after the sharding created so far when tie-breaking between candidates during explicit reshards.

Instead of source tensor sizes from unsharded tensors.

It for example prefers:

reshard lhs: {"x":(2)2}, {"y"} -> {}, {"y"}
reshard rhs: {"y"}, {} -> {"y"}, {"x"}
dot to obtain the result in sharding {}, {"x"}
all-reduce along "y"
return all-reduce

instead of:

reshard rhs: {"y"}, {} -> {"y"}, {"x":(1)2}
dot to obtain the result in sharding {"x":(2)2}, {"x":(1)2}
all-reduce along "y"
reshard to {}, {"x"}
return reshard

for the following example:

func.func @main(
%arg0: tensor<8x32xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x":(2)2}, {"y"}]>},
%arg1: tensor<32x16xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}, {}]>})
-> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {"x"}]>}) {
%0 = stablehlo.dot %arg0, %arg1
{sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"x"}]>]>}
: (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>
return %0 : tensor<8x16xf32>
}