tensorflow/CVE-2021-29516-2.patch

71 lines
3.1 KiB
Diff
Raw Normal View History

2021-09-13 10:32:00 +08:00
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