EXPECT_EQ(dltensor_out->dtype.bits, dltensor_in->dtype.bits);
  EXPECT_EQ(dltensor_out->dtype.lanes, dltensor_in->dtype.lanes);
  for (int i = 0; i < dltensor_in->ndim; ++i) {
    EXPECT_EQ(dltensor_out->shape[i], dltensor_in->shape[i]);
    if (dltensor_out->strides) {
      if (i == dltensor_in->ndim - 1) {
        EXPECT_EQ(dltensor_out->strides[i], 1);
      } else {

    std::array<TensorHandlePtr, 2> out_components;
    ExtractPerDeviceValues(context.get(), multiply_result.get(),
                           &out_components, status.get());
    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());

    ExpectScalarEq<float>(out_components[0].get(), 9.);
    ExpectScalarEq<float>(out_components[1].get(), 4.);
  }

  worker_server1.release();

    }
  }
}

int TF_OperationOutputNumConsumers(TF_Output oper_out) {
  int count = 0;
  for (const auto* edge : oper_out.oper->node.out_edges()) {
    if (edge->src_output() == oper_out.index) {
      ++count;
    }
  }
  return count;
}

int TF_OperationOutputConsumers(TF_Output oper_out, TF_Input* consumers,
                                int max_consumers) {
  int count = 0;

  TFE_DeleteExecutor(executor);
  TF_DeleteStatus(status);
}

// Read the value of variable `var` and save it into `out_value`.
void ReadVariable(TFE_Context* ctx, TFE_TensorHandle* var,
                  TFE_TensorHandle** out_value) {
  TF_Status* status = TF_NewStatus();
  TFE_Op* op = TFE_NewOp(ctx, "ReadVariableOp", status);
  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

// GraphDef type. These functions must be deleted by calling TF_DeleteFunction.
//
// If `mutate_proto_func` is non-NULL, run it over each FunctionDef proto,
// before creating a TF_Function out of the possibly mutated proto.
static std::vector<UniqueFuncPtr> CreateFunctionsFromTextProto(
    const char* text_proto,
    std::function<void(FunctionDef*)>* mutate_proto_func, TF_Status* status) {

      ExpModel, BuildGradModel(ExpModel, registry_),
      immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
}

TEST_P(CppGradients, TestMatMulGrad) {
  // TODO(vnvo2409): Figure out why `gradient_checker` does not work very
  // well with `MatMul` and remove `TestMatMul*` in
  // `mnist_gradients_test` when done.
  GTEST_SKIP();

    if (r > 0) {
      dst += r;
      offset += static_cast<uint64_t>(r);
      n -= r;  // safe as 0 < r <= n so n will never underflow
      read += r;
    } else if (r == 0) {
      TF_SetStatus(status, TF_OUT_OF_RANGE, "Read fewer bytes than requested");
      break;
    } else if (errno == EINTR || errno == EAGAIN) {
      // Retry
    } else {
      TF_SetStatusFromIOError(status, errno, posix_file->filename);
      break;

  AbstractTensorHandlePtr model_out(model_outputs[0]);

  TF_Tensor* model_out_tensor;
  TF_RETURN_IF_ERROR(GetValue(model_out.get(), &model_out_tensor));
  int num_dims_out = TF_NumDims(model_out_tensor);
  TF_DeleteTensor(model_out_tensor);

  // If the output is a scalar, then return the scalar output
  if (num_dims_out == 0) {
    outputs[0] = model_out.release();
    return absl::OkStatus();
  }

  TFE_DeleteExecutor(executor);
  TFE_ContextRemoveFunction(ctx, "AddVariablesFunction", status);
  TFE_DeleteContext(ctx);

  TF_DeleteStatus(status);

  // TODO(b/136478427): Figure out how to correctly shut the server down.
  worker_server1.release();
  worker_server2.release();
}

TEST(CAPI, TestLocalFunctionWithPackedInput) {
  TestFunctionWithPackedInput(/*remote=*/false);
}

                                             status);
  } else {
    TF_SetStatus(
        status, TF_UNIMPLEMENTED,
        absl::StrCat(
            "Trying to copy a tensor out of a parallel device. Since there "
            "are multiple components to parallel tensors, they must be "
            "unpacked explicitly.\n",
            tensorflow::unwrap(tensor)->DebugString())
            .c_str());