}];
  let arguments = (ins
    Variadic<TFTensorType>:$operands_and_static_shapes,
    DenseI32ArrayAttr:$constant_operand_indices,
    StrAttr:$metadata,
    StrAttr:$mlir_module,
    UI32Attr:$num_operands,
    DenseI32ArrayAttr:$operands_with_static_shape,
    StrAttr:$producer_name
  );

  let results = (outs
    TFTensorType:$rendezvous_key_base,
    Variadic<MlrtFutureType>:$results
  );

  }];
  let arguments = (ins
    Variadic<TF_Tensor>:$operands_and_static_shapes,
    DenseI32ArrayAttr:$constant_operand_indices,
    StrAttr:$metadata,
    StrAttr:$mlir_module,
    UI32Attr:$num_operands,
    DenseI32ArrayAttr:$operands_with_static_shape,
    StrAttr:$producer_name
  );

  let results = (outs
    TF_Tensor:$rendezvous_key_base,
    Variadic<TF_Tensor>:$results
  );

// number or a zero, it is returned as it is. If it is a negative number, it
// means it is counting backwards and will return the zero-based operand index
// for `op`.
//
// `operand_idx` should be within the range: [-num_operands, num_operands - 1].
int EvaluateOperandIdx(const int operand_idx, Operation& op) {
  if (operand_idx < 0) {
    // Calculate the actual index if a negative value is provided for
    // `operand_idx`.

      return InvalidArgument(
          "Capturing tensors from other context is not supported.");
    ret_operands.push_back(operand->getValue());
  }
  builder_.create<func::ReturnOp>(func_.getLoc(), ret_operands);

  auto arg_types = body.getArgumentTypes();
  auto result_types = body.getTerminator()->getOperandTypes();
  func_.setType(FunctionType::get(func_.getContext(), arg_types, result_types));

  // CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute
  // CHECK-SAME: constant_operand_indices = array<i32: 2>
  // CHECK-SAME: num_operands = 4
  // CHECK-SAME: operands_with_static_shape = array<i32: 0, 1, 3>

                                 std::vector<Value>& n_operands) {
  n_operands.clear();
  n_operands.reserve(non_tpu_caller->getNumOperands());
  for (auto operand : non_tpu_caller->getOperands()) {
    if (llvm::isa<BlockArgument>(operand)) {
      auto arg = llvm::cast<BlockArgument>(operand);
      n_operands.push_back(loop_operands[arg.getArgNumber()]);
      continue;
    }

  for (const auto& pos_and_input : llvm::enumerate(replicated_input_ops)) {
    auto input = pos_and_input.value();
    bool is_packed = input.getIsPacked();
    const int num_operands = input->getNumOperands();
    int num_inputs = is_packed ? 1 : num_replicas;
    if (num_operands != num_inputs)
      return input->emitOpError() << "requires " << num_inputs << " operands";
    if (is_packed) {
      packed_inputs.push_back(input->getOperand(0));

  // Ensure that the yield terminator operands matches the island results type.
  int result_count = island.getNumResults() - 1;  // -1 for the control token
  const int num_operands = yield.getNumOperands();
  if (num_operands != result_count)
    return yield.emitOpError()
           << "has " << yield.getNumOperands()
           << " operand, but island returns " << result_count;

  // Adds the "source" to the operands of the dst by creating a new dst
  // operation.
  mlir::OperationState state(dst->getLoc(), dst->getName());
  auto num_operands = dst->getNumOperands();
  state.operands.reserve(num_operands + 1);
  for (int input = 0, e = num_operands + 1; input != e; ++input) {
    if (input < dst_input) {
      state.operands.push_back(dst->getOperand(input));
    } else if (input == dst_input) {