71 lines
3.1 KiB
Diff
71 lines
3.1 KiB
Diff
From b6a0cba2b381e83a1d0a19b675ca6f7459d2d2bc Mon Sep 17 00:00:00 2001
|
|
From: Edward Loper <edloper@google.com>
|
|
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
|
|
|