Fix SIGFPE when concatenating two "empty" ragged tensors.

PiperOrigin-RevId: 221652186
This commit is contained in:
A. Unique TensorFlower 2018-11-15 10:56:29 -08:00 committed by TensorFlower Gardener
parent 7ba1642f78
commit 181b59b220
2 changed files with 9 additions and 2 deletions

View File

@ -236,8 +236,10 @@ class RaggedGatherOpBase : public OpKernel {
values_shape.set_dim(0, num_values); values_shape.set_dim(0, num_values);
TF_RETURN_IF_ERROR( TF_RETURN_IF_ERROR(
context->allocate_output(values_index, values_shape, &values_out)); context->allocate_output(values_index, values_shape, &values_out));
int64 value_size = params_dense_values_in.NumElements() / const int64 num_elements = params_dense_values_in.NumElements();
params_dense_values_in.dim_size(0); const int64 value_size =
num_elements == 0 ? 0
: (num_elements / params_dense_values_in.dim_size(0));
CallWriteValueSlices(params_dense_values_in, value_slices, value_size, CallWriteValueSlices(params_dense_values_in, value_slices, value_size,
values_out); values_out);
return ::tensorflow::Status::OK(); return ::tensorflow::Status::OK();

View File

@ -41,6 +41,11 @@ class RaggedConcatOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
] ]
@parameterized.parameters( @parameterized.parameters(
dict(
descr='Two rank-2 inputs with empty value axis=1',
rt_inputs=([[]], [[]]),
axis=1,
expected=[[]]),
dict( dict(
descr='Two rank-2 inputs (ragged_rank=1), axis=0', descr='Two rank-2 inputs (ragged_rank=1), axis=0',
rt_inputs=( rt_inputs=(