operand_not_in_launch);
  }

  // Handle pass through block arguments.
  for (OpOperand& operand :
       original_launch_op.GetBody().getTerminator()->getOpOperands()) {
    if (mlir::isa<BlockArgument>(operand.get())) {
      original_launch_op.getResult(operand.getOperandNumber())
          .replaceAllUsesWith(operand.get());
    }
  }
}

  mlir::Block& block = main.getRegion().front();
  xla::XlaBuilder builder("main");

  // Create xla_params.
  std::vector<xla::XlaOp> xla_params;
  for (mlir::BlockArgument& arg : block.getArguments()) {
    auto num = arg.getArgNumber();
    xla::Shape shape = xla::TypeToShape(arg.getType());
    xla::XlaOp argop =
        xla::Parameter(&builder, num, shape, absl::StrCat("Arg_", num));

  auto builder = mlir::OpBuilder(main_fn.getBody());
  std::vector<int> args_to_erase;
  for (int idx = 0; idx < args.size(); idx++) {
    const XlaArgument& xla_arg = args[idx];
    mlir::BlockArgument mlir_arg = main_fn.getArgument(idx);
    if (xla_arg.kind == XlaArgument::kResource) {
      mlir::Type element_type;
      if (xla_arg.type == DT_INVALID) {
        return errors::Unimplemented(absl::StrCat(

    // Additionally remove arguments from the outputs, as it provides liveness
    // throughout (functions as an invisible op above the very first op that
    // returns the arguments).
    for (const BlockArgument arg : main_func.getArguments()) {
      outputs.remove(arg);
    }
  }

  ReplaceStablehloOpsWithXlaCallModuleOp(
      SmallVector<Value>(inputs.begin(), inputs.end()),

        // Check device matching for the node defining the resource.
        if (!IsResourceMergeable(resource_attr, device_attr)) continue;
      } else {
        auto resource_arg = resource.dyn_cast<BlockArgument>();
        assert(resource_arg);
        if (resource_arg.getOwner() != &func.front()) continue;
        // Check device matching for the argument defining the resource.

    if (opr) {
      producer = opr.getOwner();
      port = {opr.getResultNumber()};
    } else {
      producer = mlir::cast<BlockArgument>(v);
      port = {0};
    }
  }
  ValuePort(PointerUnion<Operation*, BlockArgument> producer,
            SmallVector<unsigned int, 2> port)
      : producer(producer), port(port) {}

  raw_ostream& print(raw_ostream& os) const {

namespace tensorflow {
namespace tf2xla {
namespace internal {

namespace {

using llvm::ArrayRef;
using llvm::SmallVector;
using mlir::Block;
using mlir::BlockArgument;
using mlir::DenseIntElementsAttr;
using mlir::IRMapping;
using mlir::Location;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;

  // have associated tensor and buffer. Build FlatBuffer tensor and buffer for
  // other functions.
  for (unsigned i = 0, e = bb.getNumArguments(); i < e; ++i) {
    mlir::BlockArgument arg = bb.getArgument(i);
    std::string tensor_name;
    if (has_input_attr)
      tensor_name = std::string(name_mapper_.GetUniqueName(arg));
    if (tensor_name.empty()) tensor_name = absl::StrCat("arg", i);