Fix SIGFPE when concatenating two "empty" ragged tensors.
PiperOrigin-RevId: 221652186
This commit is contained in:
parent
7ba1642f78
commit
181b59b220
@ -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();
|
||||||
|
@ -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=(
|
||||||
|
Loading…
Reference in New Issue
Block a user