}

  mlir::Location loc = mlir::UnknownLoc::get(context);
  mlir::ModuleOp module = mlir::ModuleOp::create(loc);
  mlir::FunctionType func_type =
      mlir::FunctionType::get(context, input_tys, output_tys);
  llvm::StringRef func_name_str(func_name.data(), func_name.size());
  auto func = mlir::func::FuncOp::create(loc, func_name_str, func_type, {});

                             ShapedType chaining_data_type) {
  for (int i = 0; i < num_resources; ++i) {
    func.getRegion().addArgument(chaining_data_type, func.getLoc());
  }

  FunctionType ftype =
      FunctionType::get(func.getContext(), func.getBody().getArgumentTypes(),
                        func.getFunctionType().getResults());
  func.setType(ftype);
}

  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
  if (!main_fn) {
    return absl::InternalError("Could not find main function in MLIR Module.");
  }

  mlir::FunctionType func_type = main_fn.getFunctionType();
  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;
    xla::Shape xla_shape = xla::TypeToShape(input_type);

  // If there are multiple platforms, the first argument is reserved for
  // passing the platform index.
  FunctionType function_type = func_op.getFunctionType();
  ArrayRef<Type> new_input_types =
      function_type.getInputs().take_back(func_op.getNumArguments() - 1);
  func_op.setFunctionType(
      FunctionType::get(ctx, new_input_types, function_type.getResults()));
  func_op.getBody().eraseArgument(0);
}

}  // namespace

      removed_params.push_back(false);
    }
  }
  // Remove old resource-type parameters.
  callee.getBody().front().eraseArguments(removed_params);
  // Update function type.
  callee.setFunctionType(FunctionType::get(callee.getContext(),
                                           callee.getBody().getArgumentTypes(),
                                           callee.getResultTypes()));
}

    func::FuncOp func = cluster_func.getFuncOp();
    Block& block = func.front();
    for (Value read_operand : read_operands)
      block.addArgument(read_operand.getType(), loc);

    func.setType(FunctionType::get(&getContext(), block.getArgumentTypes(),
                                   func.getResultTypes()));
    cluster_func.erase();
  }
}

}  // namespace

      arg.setType(resultType);
    }
    llvm::SmallVector<Type, 8> arg_types;
    for (auto arg : func.getArguments()) arg_types.push_back(arg.getType());
    func.setType(
        FunctionType::get(func.getContext(), arg_types,
                          func.front().getTerminator()->getOperandTypes()));
    return WalkResult::advance();
  });

        absl::StrCat(iter.second.original_name.str(), ":", host.str());
    std::replace(func_name.begin(), func_name.end(), ':', '_');
    std::replace(func_name.begin(), func_name.end(), '/', '_');

    FunctionType func_type =
        FunctionType::get(context, input_types, result_types);
    Location loc = metadata.ops.front()->getLoc();
    func::FuncOp func_op = func::FuncOp::create(loc, func_name, func_type);

}

func::FuncOp CreateSessionInitFunc(ModuleOp module) {
  constexpr char kSessionInitFuncName[] = "SessionInitializerFunction";

  mlir::OpBuilder builder(module.getBodyRegion());
  auto func_type =
      FunctionType::get(module.getContext(), /*inputs=*/{}, /*results=*/{});
  auto func = builder.create<func::FuncOp>(module->getLoc(),
                                           kSessionInitFuncName, func_type);

  auto tensor_string_type =
      RankedTensorType::get({}, TF::StringType::get(builder.getContext()));
  main_func.getArguments().front().setType(tensor_string_type);
  main_func.setType(
      FunctionType::get(builder.getContext(), {tensor_string_type},
                        main_func.getFunctionType().getResults()));
  // Name of the main function for the eventual executable needs to be set.