Support mhlo.sharding attr inside backend_config of ragged_all_to…

…_all PiperOrigin-RevId: 719422910
openxla · Jan 24, 2025 · 6557723 · 6557723
1 parent 14ba309
commit 6557723
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
@@ -2275,7 +2275,7 @@ LogicalResult ExportXlaOp(CustomCallOp op, OpLoweringContext ctx) {
       auto name = attr.getName();
       return name == kCallTargetName || name == kBackendConfig ||
              name == kApiVersion || name == kCalledComputations ||
-             name == kHasSideEffect;
+             name == kHasSideEffect || name == kMhloSharding;
     };
     for (const auto& attr : op->getAttrs()) {
       if (!isSupportedAttrName(attr))