Function TEST_F

tensorflow/c/kernels_test.cc:348–385

REGISTER_OP("AllocateOutputOp1").Output("output1: float");

TEST_F(DeviceKernelOpTest, TestAllocateOutputSizeOne) {
  auto my_compute_func = [](void* kernel, TF_OpKernelContext* ctx) {
    // Allocate output
    int64_t dim = 1;
    size_t tensor_size_bytes = TF_DataTypeSize(TF_FLOAT);
    TF_Tensor* output = TF_AllocateOutput(
        /*context=*/ctx, /*index=*/0, /*dtype=*/TF_FLOAT, /*dims=*/&dim,
        /*num_dims=*/1, /*len=*/tensor_size_bytes);
    EXPECT_EQ(TF_FLOAT, TF_TensorType(output));
    EXPECT_EQ(1, TF_NumDims(output));
    EXPECT_EQ(1, TF_Dim(output, 0));

    // Set output to 3
    float* data = reinterpret_cast<float*>(TF_TensorData(output));
    float value = 3.0f;
#if GOOGLE_CUDA
    OpKernelContext* cc_ctx = reinterpret_cast<OpKernelContext*>(ctx);
    cc_ctx->eigen_gpu_device().memcpyHostToDevice(data, &value,
                                                  tensor_size_bytes);
#else
    *data = value;
#endif

    TF_Status* s = TF_NewStatus();
    TF_SetOutput(ctx, 0, output, s);
    EXPECT_EQ(TF_OK, TF_GetCode(s));

    TF_DeleteStatus(s);
    TF_DeleteTensor(output);
  };

  SetupOp("AllocateOutputOp1", "AllocateOutput1", my_compute_func);

  TF_ASSERT_OK(RunOpKernel());
  Tensor* output = GetOutput(0);
  EXPECT_EQ("Tensor<type: float shape: [1] values: 3>",
            output->DebugString(100));
}

REGISTER_OP("AllocateOutputOp0").Output("output1: float");