CreateOp op = *this;
  p << "(" << op.getInCh() << ") key("
    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") device("
    << op->getAttr("device") << ") " << op->getAttr("op_name") << "()";

  fallback_common::PrintExecuteOpCommon(p, op);
  fallback_common::PrintExecuteOpFuncAttribute(p, op);

  p << " num_args(" << op->getAttrOfType<mlir::IntegerAttr>("num_args").getInt()
    << ')';
}

    return true;
  }
  std::string shard_string;
  if (op->hasAttr(kXLAShardingAttr)) {
    shard_string =
        op->getAttrOfType<StringAttr>(kXLAShardingAttr).strref().str();
  } else {
    shard_string = op->getAttrOfType<StringAttr>(kShardingAttr).strref().str();
  }
  xla::OpSharding sharding;
  if (!shard_string.empty() && sharding.ParseFromString(shard_string)) {

// `tf_device::ClusterFuncOp`.
LogicalResult SetMetadataProtoStepMarkerLocation(
    tf_device::ClusterFuncOp op,
    tensorflow::tpu::TPUCompileMetadataProto* metadata) {
  auto step_marker_location =
      op->getAttrOfType<StringAttr>(kStepMarkerLocationAttr);
  if (!step_marker_location)
    return op.emitOpError(CreateMissingAttributeMsg(kStepMarkerLocationAttr));

  // Default to `STEP_MARK_AT_ENTRY` for step marker location if attribute is

    if (auto split = llvm::dyn_cast<TF::SplitOp>(op)) {
      if (HasDevice(split) || split->getAttrOfType<ArrayAttr>(kClassAttr))
        return WalkResult::advance();
      for (Operation* pred : IslandPredecessors(split)) {
        if (auto colocation_classes =
                pred->getAttrOfType<ArrayAttr>(kClassAttr)) {
          split->setAttr(kClassAttr, colocation_classes);
          return WalkResult::advance();

  void runOnOperation() override {
    Builder builder(&getContext());
    Dialect* tf = getContext().getLoadedDialect<TensorFlowDialect>();
    getOperation().walk([&](Operation* op) {
      if (auto device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
        // We assign default device to ops with device attribute that is empty.
        if (device_attr.getValue().empty()) {
          op->setAttr(kDeviceAttr, builder.getStringAttr(default_device_));

    // attribute modulo the number of TPU cores in the system.
    for (auto op : require_device_ordinal_ops) {
      if (op->getAttrOfType<IntegerAttr>(kDeviceOrdinalAttr).getInt() == -1) {
        OpBuilder builder(op);
        int device_ordinal =
            op->getAttrOfType<IntegerAttr>(kReplicaIdAttr).getInt() %
            device_num;
        op->setAttr(kDeviceOrdinalAttr,

namespace mlir {
namespace TF {

using ::tensorflow::kValidDeviceTypes;

LogicalResult HasValidCompilationAndReplicationAttributes(Operation& op) {
  auto replicate_attr = op.getAttrOfType<StringAttr>(kReplicationInfoAttr);
  auto compile_attr = op.getAttrOfType<StringAttr>(kCompileDeviceTypeAttr);
  if (!replicate_attr && !compile_attr) return success();
  if (!replicate_attr || !compile_attr)

      }
    } else if (StringAttr attr = arg.getDefiningOp()->getAttrOfType<StringAttr>(
                   kDeviceAttr)) {
      return op_device == attr;
    }
    // Fold tf.Identity when arg device is not defined.
    return true;
  };

  func_op.walk([&should_fold_op_func](TF::IdentityOp op) {
    StringAttr op_device = op->getAttrOfType<StringAttr>(kDeviceAttr);
    if (should_fold_op_func(op.getOperand(), op_device)) {

  double real_multiplier = scale_product / o_spec.getScale();
  output_multipliers->push_back(QuantizeMultiplier(real_multiplier));

  // output ranges
  auto min = rop->getAttrOfType<FloatAttr>("min");
  auto max = rop->getAttrOfType<FloatAttr>("max");
  output_ranges->push_back(CalculateQuantizedRange(
      o_spec.getScale(), o_spec.getZeroPoint(),
      (min ? std::optional<double>(min.getValueAsDouble()) : std::nullopt),

  StringAttr shared_name_attr = op->getAttrOfType<StringAttr>(kSharedNameAttr);
  if (shared_name_attr) {
    auto shared_name = StringRef(shared_name_attr.getValue());
    if (!shared_name.empty()) {
      return shared_name;
    }
  }
  // Attempt to parse "_class" attribute if there is no "shared_name"
  // attribute.
  ArrayAttr classes_attr = op->getAttrOfType<ArrayAttr>(kClassAttr);
  if (!classes_attr) {