Fix memory budget lower bound computation for the case of solve_nd_sharding_iteratively=true
Fix memory budget lower bound computation for the case of solve_nd_sharding_iteratively=true

** Description of the issue
When solve_nd_sharding_iteratively is turned on, say for a 2D mesh [N,M] where N >= M, we would first shard the module across the partial mesh [N,1], and then use these results to prune the search space for sharding the module across the full mesh [N,M]. An artifact of the way we have implemented this is that after sharding the module with the partial mesh, we annotate the module with the computed partial shardings (in the SetHloSharding() function).

Given that we are sharding the devices across a different number of devices each time (N vs N*M), we re-compute the memory budget lower bound (in MemoryBudgetLowerBound()). This function uses pre-existing sharding annotations, if any, on instructions for computing a lower bound on the memory consumed per device. When computing the lower bound for the case of the full mesh, we therefore end up using the partial sharding annotations on the instructions to compute the lower bound, leading to a memory budget that is ~M times larger than it should be.

** Fix
This CL fixes this issue by enforcing that preexisting shardings annotations be used only if they correspond to the same number of devices as the current device mesh.

In some cases, user annotations may not be adjusted for partial meshes. Therefore we skip the above check for instructions with user annotations. If this is found to not affect results, it can be simplified.
