// Combines the two variadic arguments ($in_tensors and $captured_tensors).
def GetBatchFunctionOpArgOperands:
    NativeCodeCall<"cast<TF::BatchFunctionOp>($0[0].getDefiningOp()).getArgOperands()">;

// Replaces `TF_BatchFunctionOp` into `TF_PartitionedCallOp` that calls the
// same $f. This may be required, for example, when inlining is desired,

      if (auto func = dyn_cast<func::FuncOp>(call.resolveCallable())) {
        PropagatePotentiallyWrittenUpFromCallee(func.getRegion(),
                                                call.getArgOperands());
      }
      return;
    }
    if (auto if_op = dyn_cast<TF::IfOp>(op)) {
      for (auto callee : {if_op.then_function(), if_op.else_function()}) {

          op.getOperation()->getOperand(to));
    }
    BitVector operands_to_erase(op->getNumOperands());
    int args_start = op->getNumOperands()
                         ? op.getArgOperands().getBase()->getOperandNumber()
                         : 0;
    operands_to_erase |= args_to_erase.lookup(func);
    operands_to_erase <<= args_start;
    op->eraseOperands(operands_to_erase);

      if (!callee_info) break;
      std::optional<int> passthrough_arg =
          callee_info.value()->GetArg(res_index);
      if (!passthrough_arg) break;
      value = call.getArgOperands()[passthrough_arg.value()];
    } else if (isa<tf_device::LaunchOp, tf_device::ClusterOp>(op)) {
      value = op->getRegion(0).front().getTerminator()->getOperand(res_index);
    } else {
      break;
    }
  }

  Region& first_region = *first->getParentRegion();
  Region& second_region = *second->getParentRegion();

  for (auto it : llvm::zip(first.getArgOperands(), second.getArgOperands())) {
    // Get the defining Op, skipping over casts.
    auto get_defining_op = [](Value value) {
      while (auto cast_op =
                 llvm::dyn_cast_or_null<CastOp>(value.getDefiningOp())) {

        }
        LLVM_DEBUG(llvm::dbgs()
                   << "Visiting call to function @" << func.getName() << "\n");
        if (failed(propagate_operands_to_callee_arguments(
                call, call.getArgOperands(), {func}, func_res)))
          return WalkResult::interrupt();
      }
      return WalkResult::advance();
    });
    if (walk_res.wasInterrupted()) return signalPassFailure();
  }
}

  let results = (outs
    Variadic<TFR_allowedResultType>:$outs);

  let extraClassDeclaration = [{
    // Get the argument operands to the called function.
    operand_range getArgOperands() { return getArgs(); }
    MutableOperandRange getArgOperandsMutable() {
      return getArgsMutable();
    }

    // Return the callee of this operation.

  );

  let results = (outs
    Variadic<AnyType>:$results
  );

  let extraClassDeclaration = [{
    // Gets the argument operands to the called function.
    operand_range getArgOperands() { return getArgs(); }
    MutableOperandRange getArgOperandsMutable() {
      return getArgsMutable();
    }
    // returns the function that this operation will launch.
    func::FuncOp getFuncOp() {

  );

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

  let extraClassDeclaration = [{
    // Gets the argument operands to the called function.
    operand_range getArgOperands() { return getArgs(); }
    MutableOperandRange getArgOperandsMutable() {
      return getArgsMutable();
    }

    // Returns the callee of this operation.

        auto callee = dyn_cast<func::FuncOp>(call.resolveCallable());
        for (const auto& entry : AccessedGradients({callee}, module))
          for (const string& source : entry.getSecond())
            insert(call.getArgOperands()[entry.getFirst()], source, func_block);
      }
    });
  }
  return result;
}

// Contains cached information for decomposed callee functions for (stateful)
// partitioned call ops.