From c988ae6d70bff223ae7be8ea02368d968e431124 Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Fri, 27 Sep 2024 11:58:23 -0700
Subject: [PATCH] In a previous change, we throw an error if we encounter an
 unknown sharding 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: 679676685
---
 xla/hlo/experimental/auto_sharding/auto_sharding.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/xla/hlo/experimental/auto_sharding/auto_sharding.cc b/xla/hlo/experimental/auto_sharding/auto_sharding.cc
index d69764435604cc..29ad335d888e23 100644
--- a/xla/hlo/experimental/auto_sharding/auto_sharding.cc
+++ b/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");