// updated in the function.
  llvm::SmallSet<int, 4> indexes;
  for (BlockArgument &arg : func.getArguments()) {
    if (tensor_list_args.contains(arg.getArgNumber())) {
      for (const mlir::OpOperand &use : arg.getUses()) {
        mlir::Operation *op = use.getOwner();
        // Currently we only check if the tensorlist argument is consumed by
        // `TensorListPushBack` or `TensorListResize`, since those are the only

Value GetResultInBlock(Value value, llvm::DenseSet<Block*> target_blocks) {
  while (target_blocks.find(value.getParentBlock()) == target_blocks.end()) {
    Value new_value = nullptr;
    for (OpOperand& operand : value.getUses()) {
      Operation* owner = operand.getOwner();
      if (!llvm::isa<tf_device::ReturnOp>(owner)) continue;
      Operation* parent = owner->getParentOp();  // op that owns the "return"

      }
      return_operands.push_back(resource.live_value);
      result_types.push_back(resource.live_value.getType());
    }

    if (is_var_handle && !resource.read) {
      assert(block.getArgument(index).getUses().empty());
      argument_indices_to_remove.push_back(index);
    } else {
      if (is_var_handle) {
        // Add resource_name attribute to VarHandleOp read.
        function.setArgAttr(

    for (Operation* user : island.getControl().getUsers()) {
      DCHECK_EQ(user->getParentOp(), graph);
      try_update_current_candidate(user);
    }

    // Check island data results.
    Block& graph_body = llvm::cast<GraphOp>(graph).GetBody();
    for (Value result : island.getOutputs()) {
      for (Operation* user : result.getUsers()) {
        Operation* def = graph_body.findAncestorOpInBlock(*user);

  // GraphOp, YieldOp and NextIterationSourceOp don't have control inputs so
  // exclude them below.
  for (Value out : island_op.getOutputs()) {
    for (auto& use : out.getUses()) {
      Operation* owner = use.getOwner();
      if (owner->getDialect() == island_op->getDialect() &&
          !llvm::isa<tf_executor::GraphOp, tf_executor::YieldOp,

    if (constant_op.getValue().getNumElements() > 32) continue;
    while (!constant_op.getResult().hasOneUse()) {
      auto new_constant_op = builder.clone(*constant_op.getOperation());
      constant_op.getResult().getUses().begin()->assign(
          dyn_cast<mlir::stablehlo::ConstantOp>(new_constant_op));
    }
  }
}

void ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass::
    runOnOperation() {

    if (op.getNumOperands() != 0 || op.getNumResults() != 1 ||
        !HasSingleOpInBlock<YieldOp>(&op.GetBody()))
      return failure();

    for (auto &use : llvm::make_early_inc_range(op.getControl().getUses()))
      use.getOwner()->eraseOperand(use.getOperandNumber());

    rewriter.eraseOp(op);

    return success();
  }
};

          nullptr};
    }
    // Check the current operation has a consumer node.
    Operation* consumer_op =
        current->getResult(0).getUses().begin()->getOwner();
    if (!consumer_op) {
      return {
          rewriter.notifyMatchFailure(current, "op doesn't have consumer node"),
          nullptr};
    }

    for (const auto& idx_result : llvm::enumerate(island.getOutputs())) {
      Value result = idx_result.value();

      bool has_external_use = false;
      for (OpOperand& use : llvm::make_early_inc_range(result.getUses())) {
        if (wrapped_ops.count(use.getOwner()))
          use.set(yield_op.getOperand(idx_result.index()));
        else
          has_external_use = true;
      }
      if (has_external_use) {

llvm::SmallVector<Operation*> GetSuccessors(Operation* op) {
  llvm::SmallVector<Operation*> successors;
  for (auto result : op->getParentOp()->getOpResults()) {
    for (auto& use : result.getUses()) {
      auto succ = use.getOwner();
      successors.push_back(succ);
    }
  }
  return successors;
}
// Gets the predecessors of an op wrapped in tf_executor.island.