experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix TF_AllocateOutput to actually call TF_SetOutput.

Browse Source 
Before this change, TF_AllocateOutput would simply allocate a tensor but not
place that tensor in the specified output position. This change necessitates
the addition of a TF_Status* output parameter.

PiperOrigin-RevId: 269403835
This commit is contained in:
James Ring 2019-09-16 13:12:34 -07:00 committed by TensorFlower Gardener
parent 98b54f7d93
commit 9cd757ab16
3 changed files with 20 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/c/kernels.cc
View File

@ -244,11 +244,18 @@ int64_t TF_StepId(TF_OpKernelContext* ctx) {
TF_Tensor* TF_AllocateOutput(TF_OpKernelContext* context, int index,
TF_DataType dtype, int64_t* dims, int num_dims,
size_t len) {
size_t len, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(context);
tensorflow::AllocatorAttributes attr = cc_ctx->output_alloc_attr(index);
auto* allocator = cc_ctx->get_allocator(attr);
void* data = tensorflow::allocate_tensor("TF_AllocateOutput", len, allocator);
return TF_NewTensor(dtype, dims, num_dims, data, len,
TF_Tensor* result = TF_NewTensor(dtype, dims, num_dims, data, len,
tensorflow::deallocate_buffer, allocator);
TF_SetOutput(context, index, result, status);
if (TF_GetCode(status) != TF_OK) {
TF_DeleteTensor(result);
return nullptr;
}
return result;
}

2
tensorflow/c/kernels.h
View File

@ -188,7 +188,7 @@ TF_CAPI_EXPORT extern void TF_OpKernelConstruction_GetAttrInt32(
TF_CAPI_EXPORT TF_Tensor* TF_AllocateOutput(TF_OpKernelContext* context,
int index, TF_DataType dtype,
int64_t* dims, int num_dims,
size_t len);
size_t len, TF_Status* status);
#ifdef __cplusplus
} /* end extern "C" */

24
tensorflow/c/kernels_test.cc
View File

@ -348,11 +348,13 @@ REGISTER_OP("AllocateOutputOp1").Output("output1: float");
TEST_F(DeviceKernelOpTest, TestAllocateOutputSizeOne) {
auto my_compute_func = [](void* kernel, TF_OpKernelContext* ctx) {
// Allocate output
TF_Status* s = TF_NewStatus();
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);
/*num_dims=*/1, /*len=*/tensor_size_bytes, s);
EXPECT_EQ(TF_OK, TF_GetCode(s));
EXPECT_EQ(TF_FLOAT, TF_TensorType(output));
EXPECT_EQ(1, TF_NumDims(output));
EXPECT_EQ(1, TF_Dim(output, 0));
@ -368,10 +370,6 @@ TEST_F(DeviceKernelOpTest, TestAllocateOutputSizeOne) {
*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);
};
@ -388,20 +386,18 @@ REGISTER_OP("AllocateOutputOp0").Output("output1: float");
TEST_F(DeviceKernelOpTest, TestAllocateEmptyOutput) {
auto my_compute_func = [](void* kernel, TF_OpKernelContext* ctx) {
TF_Status* s = TF_NewStatus();
// Allocate empty output
int64_t dim = 0;
TF_Tensor* output = TF_AllocateOutput(
/*context=*/ctx, /*index=*/0, /*dtype=*/TF_FLOAT, /*dims=*/&dim,
/*num_dims=*/1, /*len=*/0);
/*num_dims=*/1, /*len=*/0, s);
EXPECT_EQ(TF_OK, TF_GetCode(s));
EXPECT_EQ(TF_FLOAT, TF_TensorType(output));
EXPECT_EQ(1, TF_NumDims(output));
EXPECT_EQ(0, TF_Dim(output, 0));
TF_Status* s = TF_NewStatus();
TF_SetOutput(ctx, 0, output, s);
EXPECT_EQ(TF_OK, TF_GetCode(s));
TF_DeleteStatus(s);
TF_DeleteTensor(output);
};
@ -418,12 +414,14 @@ REGISTER_OP("AllocateOutputOp2x3").Output("output1: float");
TEST_F(DeviceKernelOpTest, TestAllocateOutputSize2x3) {
auto my_compute_func = [](void* kernel, TF_OpKernelContext* ctx) {
TF_Status* s = TF_NewStatus();
// Allocate 2x3 output
int64_t dim[2] = {2, 3};
size_t tensor_size_bytes = 6 * TF_DataTypeSize(TF_FLOAT);
TF_Tensor* output = TF_AllocateOutput(
/*context=*/ctx, /*index=*/0, /*dtype=*/TF_FLOAT, /*dims=*/dim,
/*num_dims=*/2, /*len=*/tensor_size_bytes);
/*num_dims=*/2, /*len=*/tensor_size_bytes, s);
EXPECT_EQ(TF_OK, TF_GetCode(s));
EXPECT_EQ(TF_FLOAT, TF_TensorType(output));
EXPECT_EQ(2, TF_NumDims(output));
EXPECT_EQ(2, TF_Dim(output, 0));
@ -440,10 +438,6 @@ TEST_F(DeviceKernelOpTest, TestAllocateOutputSize2x3) {
memcpy(data, value, tensor_size_bytes);
#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);
};