tensorflow::DeviceNameUtils::ParsedName name;
      if (tensorflow::DeviceNameUtils::ParseFullName(device.str(), &name)) {
        if (name.type == "TPU_REPLICATED_CORE") {
          // TODO(hanxiongwang): Add check for out of bound of name.id
          return llvm::StringRef(logical_device_vec[name.id]);
        }
      }
    }
  }
  return std::nullopt;
}

      } else if (quant_specs_.inference_type == tensorflow::DT_HALF) {
        quantizeOpAsFloat16(rewriter, op, quant_op);
        quantized = true;
      }
    }
    return quantized;
  }

  // Add asymmetric input quantization attribute. MLIR dynamic quantization
  // supports only the case that the value of the attribute equals to true. For
  // details, see tensorflow/compiler/mlir/lite/quantization/quantization.td

                                    impl()->colocation_constraints_.end());
    // Sort the set.
    std::sort(constraints.begin(), constraints.end());
    // Add loc:@ prefix
    std::transform(constraints.begin(), constraints.end(), constraints.begin(),
                   [](const string& s) {
                     return strings::StrCat(kColocationGroupPrefix, s);

      if (memory_bandwidth >= 0) {
        builder.set_memory_bandwidth(memory_bandwidth);
      }
    }
    // TODO(annarev): Add gflops field in DeviceDescription and set it here.
    // TODO(annarev): Perhaps add `supports_unified_memory` in
    // DeviceDescription.
    return builder.Build();
  }

  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override {

                       graph_op_def.deprecation().version(), ":\n",
                       graph_op_def.deprecation().explanation(), ".\n");
  } else if (api_def.summary().empty()) {
    comment = "TODO: add doc.\n";
  } else {
    comment = strings::StrCat(api_def.summary(), "\n");
  }
  if (!api_def.description().empty()) {
    strings::StrAppend(&comment, "\n", api_def.description(), "\n");
  }

  OperationState new_state(loc, tf_op_name, inputs, output_types, attr_list);
  Operation* new_op = rewriter.create(new_state);
  if (materialize_derived_attrs_) {
    for (const auto& attr : derived_attrs) {
      // Add or update the derived attribute with the value. Skip the fixed
      // element type attributes, in case they are present in the NodeDef.
      if (!fixed_elt_type_attrs_.contains(attr.first())) {

};

void DeferActivationTransposePass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<DeferActivationTransposeForAddOp,
               DeferActivationTransposeForMaxPoolReduceWindowOp,
               DeferActivationTransposeForMaxOp>(&ctx);
  if (failed(applyPatternsAndFoldGreedily(func_op, std::move(patterns)))) {

void PrepareQuantizeDRQPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  ModuleOp module_op = getOperation();

  populateWithGenerated(patterns);
  patterns.add<PrepareDRQQuantizableOp>(ctx, quant_specs_, op_set_,
                                        enable_per_channel_quantization_);
  FrozenRewritePatternSet frozen_patterns(std::move(patterns));

    const auto& api_def_attr(op_info.api_def.attr(i));
    // Skip inferred arguments
    if (op_info.inferred_input_attrs.count(attr.name()) > 0) continue;
    // Skip if it has default value (TODO(unda, b/249345399): add our custom
    // values)
    if (api_def_attr.has_default_value()) continue;
    // TODO(unda, b/253432797): handle unimplemented input attribute types

                               OpType& found_op) {
  // Find the TPUReplicationMetadata or TPUCompilationResult ops which will be
  // cloned/inserted into each region. We add them to the merged_set so that
  // they're ignored when extracting the four main functions.
  found_op = nullptr;
  for (OpType op : func.getOps<OpType>()) {
    if (found_op != nullptr) {