block_arg.getLoc(),
      DenseIntElementsAttr::get(
          RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
          group_assignment_val));

  StringAttr reduce_op = builder.getStringAttr("Add");
  StringAttr mode = builder.getStringAttr("CrossReplica");
  return builder.create<XlaAllReduceOp>(block_arg.getLoc(), block_arg.getType(),

        parsed_name.id = 0;
        parsed_name.has_id = true;
      }
      output.push_back(mlir::NamedAttribute(
          mlir::StringAttr::get(context, "device"),
          mlir::StringAttr::get(
              context,
              tensorflow::DeviceNameUtils::ParsedNameToString(parsed_name))));
    } else {
      output.push_back(attr);
    }
  }
  return mlir::success();
}

class TFL_AnyStrAttrOf<list<string> cases> : StringBasedAttr<
  CPred<!foldl(
      "$_self.cast<StringAttr>().getValue() == \"" # !head(cases) # "\"",
      !foreach(case, !tail(cases),
               "$_self.cast<StringAttr>().getValue() == \"" # case # "\""),
      prev, cur, prev # " || " # cur)>,
  "string attribute whose value is " #

LogicalResult AssignDevicesInRegion(const Dialect* tf_dialect,
                                    tf_device::LaunchOp launch,
                                    Region& region) {
  auto parallel_group_attr =
      launch->getAttrOfType<StringAttr>(TF::kParallelExecAnnotation);
  auto result = region.walk([&](Operation* op) -> WalkResult {
    if (op->getDialect() != tf_dialect) return WalkResult::advance();

    if (parallel_group_attr) {

namespace internal {

namespace {

using mlir::BlockArgument;
using mlir::failure;
using mlir::LogicalResult;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;
using mlir::StringAttr;
using mlir::success;
using mlir::Value;
using mlir::WalkResult;
using mlir::func::FuncOp;
using mlir::TF::ReadVariableOp;
using mlir::tf_device::ReplicateOp;

#define GEN_PASS_DEF_HOISTBROADCASTREADPASS

  return op->hasAttr(TF::kReplicationInfoAttr);
}

StringAttr GetReplicationAttr(mlir::Operation* op) {
  return op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
}

StringAttr GetReplicationAttr(TF::TPUCompilationResultOp op) {
  // Special case for getting the replication region for
  // TPUCompilationResultsOp.
  return op->getAttrOfType<StringAttr>(kTpuCompilationStatus);
}

  Operation* op = val.getDefiningOp();
  // val should always have a defining op because cluster inputs always have
  // defining ops.
  assert(op);
  StringAttr sharding_attr = op->getAttrOfType<StringAttr>("_XlaSharding");
  if (!sharding_attr)
    return context_op->emitOpError()
           << "A map_outside_compilation op's input should have an explicit "

  DeviceNameUtils::ParsedName device;

  for (const auto& kv : llvm::enumerate(array_attr)) {
    const int idx = kv.index();

    auto string_attr = mlir::dyn_cast<mlir::StringAttr>(kv.value());
    if (!string_attr)
      return op->emitOpError(llvm::formatv(
          "bad '{0}' attribute at index {1}, not a string", kDevicesAttr, idx));

      ConversionPatternRewriter& rewriter) const final;
};

// Ops that have a call_target_name starting with the prefix "custom_call." and
// backend_config of type StringAttr (if specified) should be legalized (i.e.
// considered to be illegal if still present after this pass is complete).
std::optional<bool> IsCustomCallLegal(mhlo::CustomCallOp op);

}  // namespace odml
}  // namespace mlir

    auto replication_info_attr_str =
        mlir::dyn_cast<StringAttr>(replication_info_attr);
    if (!replication_info_attr_str ||
        replication_info_attr_str.getValue().empty())
      return metadata_op.emitError() << kBadReplicateInfoAttrMsg;

    // Remove `name` attribute.
    attrs.erase(StringAttr::get(metadata_op.getContext(), kNameAttr));