double GetOpCost(mlir::Operation* op) const override {
    int64_t count;
    if (ArithmeticCountUtilHelper::GetFirstOutputCount(op, &count))
      return kGPUArithmeticUnitCost * count;
    return kGPUDefaultFixedValuedCost;
  }

  bool IsOpSupported(mlir::Operation* op) const override {
    InferenceType inference_type = GetInferenceType(op);
    if (inference_type != FLOAT) {
      return false;

    func.walk([&](mlir::TF::ReadVariableOp op) {
      auto sym = lookupGlobalTensor(func, op.getResource(), syms, opToName);
      success &= !!sym;
      if (!success) return;
      OpBuilder builder(op);
      auto load = builder.create<mlir::ml_program::GlobalLoadOp>(
          op.getLoc(), op.getValue().getType(), sym);
      op.getValue().replaceAllUsesWith(load.getResult());
      op.erase();
    });

  auto update_from_op = [&](Operation* op) {
    been_defined.insert(op->getResults().begin(), op->getResults().end());
    for (Value input : op->getOperands()) {
      if (been_defined.contains(input)) {
        continue;
      }
      results.insert(input);
    }
  };
  for (Operation* op : partition_ops) {
    update_from_op(op);
    op->walk<WalkOrder::PreOrder>([&](Block* nested_block) {

    // Create a new dequantize op that is propagated.
    rewriter.setInsertionPointAfter(user_op);
    TF::PartitionedCallOp new_dequantize_op =
        cast<TF::PartitionedCallOp>(rewriter.clone(*original_dequantize_op));

    // Skip the original dequant op and connect the op before dequantize to the
    // user op.
    user_op->setOperand(user_idx, op_before_dequantize);

      // old ConstantOp is guaranteed to have one F32->F16 convert op regardless
      // of its number of users.
      rewriter.setInsertionPointAfter(op);
      // create new F16 constant op in that location
      ConstantOp new_const = rewriter.create<ConstantOp>(
          op->getLoc(), new_result_type, new_value_attr);
      ConvertOp dcast =
          rewriter.create<ConvertOp>(op->getLoc(), old_result_type, new_const);

      })
      .wasInterrupted();
}

// Returns whether `op` or any ancestors of `op` are outside compiled.
bool HasOutsideCompilationAncestor(Operation* op) {
  while (op) {
    if (op->hasAttr(kXlaOutsideCompilationAttr)) {
      return true;
    }
    op = op->getParentOp();
  }
  return false;
}

// Returns whether any ancestors of `op` are outside compiled.

}

// Checks if an op is a TF/XLA communication op.
bool IsCommunicationOp(Operation* op) {
  return isa<TF::_XlaHostComputeMlirOp, TF::XlaSendToHostOp,
             TF::XlaRecvFromHostOp>(op);
}

// Checks if an op is a supported HLO control flow op.
bool IsControlFlowOp(Operation* op) { return isa<IfOp, WhileOp>(op); }

// Collects control flow op ancestors of a given op, up until FuncOp. If any

        };

    ValueRange opInput = op.getInput();
    TypeRange scf_block_arguments_type = opInput.getType();

    // Create the `scf.while` op.
    auto scf_while_op = rewriter.create<scf::WhileOp>(
        op.getLoc(), op.getResultTypes(), opInput);

    // Create the `before` block of the `scf.while` op (with an `scf.condition`
    // op as the terminator). Note that the arguments' type of this block is

    auto calledComputations = op.getCalledComputations();
    if (!calledComputations || calledComputations.size() != 1)
      return op->emitError("expected exactly one called_computation");

    auto decomposition = mlir::cast<FlatSymbolRefAttr>(calledComputations[0]);

    auto composite = rewriter.create<mlir::stablehlo::CompositeOp>(
        op.getLoc(), op.getResultTypes(), op.getOperands(), name.str(), attrs,
        decomposition.getValue());

        llvm::all_of(cluster_op.getUsers(), [&](Operation* op) {
          return op == terminator || tail_outside_compiled_ops_set.count(op);
        });
    if (!can_be_extracted) continue;

    // Collect operands of cluster op that are generated within the cluster.
    // These values should be returned by the cluster.
    cluster_op.walk([&](Operation* op) {
      for (Value operand : op->getOperands()) {