}

      rewriter.setInsertionPoint(cloned_func.getBody().front().getTerminator());
      rewriter.create<TF::_XlaSendFromHostOp>(
          func.getLoc(),
          cloned_func.getBody().front().getTerminator()->getOperands(),
          /*dynamic_key=*/dynamic_key, op.getRecvKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),
          rewriter.getStringAttr("TPU"));
    }

// The current implementation is based on the number of operands.
static bool PreferResultScale(Operation* op) {
  int float_operands = 0;
  for (auto operand : op->getOperands()) {
    if (auto operand_type = dyn_cast<ShapedType>(operand.getType())) {
      if (isa<FloatType>(operand_type.getElementType())) {
        if (++float_operands > 1) return true;
      }
    }
  }

  //
  // Sharding configurations are added to the tf_device.ClusterFunc as an
  // attribute and the function as an argument attribute.
  for (auto operand_and_arg :
       llvm::zip(cluster_func.getOperands(), function_block.getArguments())) {
    Value operand = std::get<0>(operand_and_arg);
    BlockArgument arg = std::get<1>(operand_and_arg);

    if (auto operand_sharding = GetXlaShardingFromOperand(operand)) {

  if (num_cores_per_replica != 1) return success();

  llvm::SetVector<Value> bcasts;
  cluster->walk([&](Operation* op) {
    if (op == cluster) return WalkResult::advance();
    for (auto operand : op->getOperands()) {
      Operation* scope = operand.getParentBlock()->getParentOp();
      if (scope->isProperAncestor(replicate)) {
        bcasts.insert(operand);
      }
    }
    return WalkResult::advance();
  });

    auto walk_res = func_op.walk([&](Operation* op) {
      if (auto while_op = dyn_cast<WhileOp>(op)) {
        if (failed(propagate_operands_to_callee_arguments(
                while_op, while_op.getOperands(),
                {while_op.body_function(), while_op.cond_function()},
                func_res)))
          return WalkResult::interrupt();
      } else if (auto if_op = dyn_cast<IfOp>(op)) {

  std::string custom_option_buffer;
  CreateFlexOpCustomOptions(op_name, node_def_str, custom_option_buffer);
  auto flex_op = builder.create<TFL::CustomOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(), flex_op_name,
      CustomOptionForFlexOp(&builder, custom_option_buffer));
  op->replaceAllUsesWith(flex_op);
  op->erase();
  return true;
}

// Sets the "no_fallback" attribute.

    : region_(&region) {
  if (region.empty()) return;

  assert(llvm::hasSingleElement(region.getBlocks()));

  auto results = region.front().getTerminator()->getOperands();
  if (results.empty()) return;

  backtracked_values_.reserve(results.size());
  for (auto result : results)
    backtracked_values_.push_back(backtrack_analysis.BacktrackValue(result));
}

  // UniformQuantized Ops are considered legal if its qint operands and
  // results are connected to TF CastOp.
  return op && llvm::all_of(op->getResults(), IsQintValueQintToIntCast) &&
         llvm::all_of(op->getOperands(), IsQintValueDefinedByIntToQintCast);
}

bool IsCastOpLegal(TF::CastOp cast_op) {
  // Consider qint <-> qint casts illegal.
  if (IsIllegalType(cast_op.getSrcT()) && IsIllegalType(cast_op.getDstT())) {

    func_metadata.original_name = func_op.getName();
    func_metadata.insertion_point = ++Block::iterator(func_op);
    func_metadata.ops.push_back(op);

    for (Value value : op->getOperands()) {
      std::string value_device = "";

      // If the value is defined as an argument of the func_op, adds it to
      // the argument list of the function that uses this op.

      }

      write_op.erase();
    }
  }

  const int64_t num_results_before = function.getNumResults();
  auto return_operands = llvm::to_vector<4>(return_op.getOperands());
  auto result_types = llvm::to_vector<4>(return_op.getOperandTypes());
  llvm::SmallVector<std::pair<int64_t, llvm::StringRef>, 4>
      output_only_resources;