"Encountered while processing input ", i, " into function '", fn_name,
        "'");
    TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNonRefOutput(node, idx),
                                    "Encountered while processing input ", i,
                                    " into function '", fn_name, "'");

    input_tensors->emplace_back(node, idx);

    const auto& iter = input_nodes->find(node);
    if (iter == input_nodes->end()) {

  implicitly_broadcast_tensors.reserve(inputs.size());  // not tight
  for (const auto& input : inputs) {
    if (absl::holds_alternative<TFE_TensorHandle*>(input)) {
      if (operation_name == std::string("_EagerConst")) {
        // Non-parallel tensors from _EagerConst/tf.constant are implicitly
        // broadcast, i.e. set as the input to each parallel operation. This

    TFE_TensorHandle* input = TFE_OpGetFlatInput(original_op, j, s);
    if (TF_GetCode(s) != TF_OK) return;
    const char* input_device = TFE_TensorHandleDeviceName(input, s);
    if (TF_GetCode(s) != TF_OK) return;
    if (dev->device_name == input_device) {
      LoggedTensor* t = reinterpret_cast<LoggedTensor*>(
          TFE_TensorHandleDevicePointer(input, s));
      if (TF_GetCode(s) != TF_OK) return;

      "      op: 'Add'"
      "      input: 'read0:value:0'"
      "      input: 'read1:value:0'"
      "      attr {"
      "        key: 'T'"
      "        value {"
      "          type: DT_FLOAT"
      "        }"
      "      }"
      "    }"
      "    node_def {"
      "      name: 'add2'"
      "      op: 'Add'"
      "      input: 'add1:z:0'"
      "      input: 'read2:value:0'"
      "      attr {"

    TFE_OpReset(op_.get(), operation_name, device_.c_str(), status);
    if (TF_GetCode(status) != TF_OK) return;
  }
  TFE_OpAddAttrs(op_.get(), attributes);
  for (int input_index = 0; input_index < inputs.size(); ++input_index) {
    TFE_OpAddInput(op_.get(), inputs[input_index], status);
    if (TF_GetCode(status) != TF_OK) return;
  }
  std::vector<TFE_TensorHandle*> unwrapped_results(expected_max_outputs);

                         int input_index, bool use_function,
                         AbstractTensorHandle** numerical_grad) {
  vector<AbstractTensorHandle*> theta_inputs(inputs.size());
  for (int i{}; i < inputs.size(); ++i) {
    theta_inputs[i] = inputs[i];
  }

  AbstractTensorHandle* theta =
      theta_inputs[input_index];  // parameter we are grad checking

  std::unique_ptr<SingleMachine> cluster(new SingleMachine(5 * 60, 3, 0));
  TF_ASSERT_OK(cluster->Provision());

  TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
                                          cluster->GetDeviceNames());
  GrapplerItem item;
  CHECK(fake_input.NextItem(&item));

  TF_Status* status = TF_NewStatus();
  TF_GraphProperties* graph_properties =

    TFE_Context* context, TFE_TensorHandle* input,
    std::array<TensorHandlePtr, num_replicas>* components, TF_Status* status) {
  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
      TFE_NewOp(context, "TPUReplicatedOutput", status), TFE_DeleteOp);
  if (TF_GetCode(status) != TF_OK) return;
  TFE_OpSetAttrInt(op.get(), "num_replicas", num_replicas);
  TFE_OpAddInput(op.get(), input, status);
  if (TF_GetCode(status) != TF_OK) return;

}

TensorHandlePtr CollectiveSum(TFE_Context* context, TFE_TensorHandle* input,
                              int group_size, TF_Status* status) {
  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
      TFE_NewOp(context, "CollectiveReduce", status), TFE_DeleteOp);
  if (TF_GetCode(status) != TF_OK) return nullptr;

  const char* device = TFE_TensorHandleDeviceName(input, status);
  if (TF_GetCode(status) != TF_OK) return nullptr;

  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
  const tensorflow::Tensor& arg_tensor = cc_ctx->input(index);
  absl::Status cc_status;
  if (arg_tensor.dtype() != tensorflow::DT_RESOURCE) {
    cc_status = absl::InvalidArgumentError(
        absl::StrCat("Trying to obtain resource handle from Input[", index,
                     "], which is not type DT_RESOURCE."));
    status->status = cc_status;
    return nullptr;