From b6a0cba2b381e83a1d0a19b675ca6f7459d2d2bc Mon Sep 17 00:00:00 2001 From: Edward Loper Date: Tue, 25 Aug 2020 08:12:53 -0700 Subject: [PATCH 1/1] Fix segmentation fault in tf.map_fn when fn_output_spec is a RaggedTensorSpec and the input tensor has shape [0, ...]. PiperOrigin-RevId: 328332518 Change-Id: I6aff03152bbc96507fb6c5f89b05722f3cc30164 --- .../kernels/ragged_tensor_from_variant_op.cc | 16 +++++++++++++++- .../python/ops/ragged/ragged_map_fn_op_test.py | 15 +++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc index ad0712e6fd0..aa736ad7f60 100644 --- a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc +++ b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc @@ -175,8 +175,22 @@ Status NestedStackRaggedTensors( } } + // If the variant tensor input is empty, then we have no way to determine + // the correct shape for the dense_values. (It must have rank>=1, and its + // outer dimension must be 0, but we don't know its shape beyond that.) + // For now, we just use a shape of `[0]` in this case. + // TODO(edloper): Update this op with an attribute containing information + // about dense_values shape. If it's `None`, then we'll probably still have + // to use shape=[0] here, but if we have more info, then we can use it. + // E.g., in map_fn, we may have shape info from the RaggedTensorSpec. + TensorShape component_values_shape; + if (ragged_components.empty()) { + component_values_shape = TensorShape({0}); + } else { + component_values_shape = ragged_components[0].values.shape(); + } + // Populate values. - TensorShape component_values_shape = ragged_components[0].values.shape(); int values_size = component_values_shape.dim_size(0); for (int i = 1; i < ragged_components.size(); i++) { if (ragged_components[i].values.dims() != component_values_shape.dims()) { diff --git a/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py b/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py index 8a40e396a68..bead4923a0a 100644 --- a/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py +++ b/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py @@ -150,6 +150,21 @@ class RaggedMapOpTest(test_util.TensorFlowTestCase, result_dtype=ragged_tensor.RaggedTensorType( dtype=dtypes.int64, ragged_rank=4), ), + # [d1] -> [d1, (d2), (d3)] + dict( + fn=ragged_math_ops.range, + elems=np.array([1, 2, 3], np.int64), + expected_output=[[[0]], [[0, 1]], [[0, 1, 2]]], + result_dtype=ragged_tensor.RaggedTensorType( + dtype=dtypes.int64, ragged_rank=2)), + # [0] -> [0, (d2), (d3)] (github issue #36232) + dict( + fn=ragged_math_ops.range, + elems=np.zeros([0], np.int64), + expected_output=[], + expected_ragged_rank=2, + result_dtype=ragged_tensor.RaggedTensorType( + dtype=dtypes.int64, ragged_rank=2)), ]) def testRaggedMap( -- 2.27.0