In a previous change, we throw an error if we encounter an unknown sh…

…arding when saving shardings for instructions. However, this ingored the fact that we deliberately replace some module parameter/root shardings with unknown sharding objects. This CL makes the condition tigher so we only throw an error when we encounter unknown sharding objects intended for shard_as or shard_like annotations, which was the original intention anyway.

PiperOrigin-RevId: 679734540

xla/hlo/experimental/auto_sharding/auto_sharding.cc

@@ -2572,7 +2572,8 @@ absl::Status SaveShardingForInstruction(
     if (!inst->has_sharding()) {
       return absl::OkStatus();
     }
-    if (inst->sharding().IsUnknown()) {
+    if (inst->sharding().IsUnknown() &&
+        (inst->sharding().IsShardLike() || inst->sharding().IsShardAs())) {
       return absl::UnimplementedError(
           "Auto-sharding currently does not support shard_as/shard_like "
           "sharding annotations");