packages/tensorflow
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!28 fix CVE-2021-29516 CVE-2021-37679

Browse Source 
From: @starlet-dx
Reviewed-by: @zhengzhenyu
Signed-off-by: @zhengzhenyu
This commit is contained in:
openeuler-ci-bot 2021-09-14 01:25:36 +00:00 committed by Gitee
commit 08617492e8
6 changed files with 1300 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

199
CVE-2021-29516-1.patch Normal file
View File

@ -0,0 +1,199 @@
From ce47a396ff795bdb6cf48eb53dbcba46cb51fa7d Mon Sep 17 00:00:00 2001
From: Katherine Tian <kattian@google.com>
Date: Tue, 30 Jun 2020 04:12:11 +0000
Subject: [PATCH 1/1] TensorKey class and TensorMap tests
---
tensorflow/core/BUILD | 1 +
tensorflow/core/framework/BUILD | 70 ++++++++++++++++++++++++++
tensorflow/core/framework/tensor_key.h | 64 +++++++++++++++++++++++
tensorflow/core/kernels/BUILD | 1 +
4 files changed, 136 insertions(+)
create mode 100644 tensorflow/core/framework/tensor_key.h
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index d0be6ee9..6e745b4e 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -495,6 +495,7 @@ tf_cuda_library(
"//tensorflow/core/framework:shared_ptr_variant.h",
"//tensorflow/core/framework:stats_aggregator.h",
"//tensorflow/core/framework:tensor.h",
+ "//tensorflow/core/framework:tensor_key.h",
"//tensorflow/core/framework:tensor_shape.h",
"//tensorflow/core/framework:tensor_slice.h",
"//tensorflow/core/framework:tensor_types.h",
diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 9b6ddb2a..093f0545 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -209,6 +209,7 @@ filegroup(
"shared_ptr_variant.h",
"stats_aggregator.h",
"tensor.h",
+ "tensor_key.h",
"tensor_reference.h",
"tensor_shape.h",
"tensor_slice.h",
@@ -760,6 +761,75 @@ tf_cuda_library(
alwayslink = 1,
)
+tf_cuda_library(
+ name = "tensor_key",
+ srcs = [
+ "log_memory.cc",
+ "tensor.cc",
+ "typed_allocator.cc",
+ "types.cc",
+ "variant.cc",
+ "variant_op_registry.cc",
+ "variant_tensor_data.cc",
+ ],
+ hdrs = [
+ "log_memory.h",
+ "register_types.h",
+ "tensor.h",
+ "tensor_key.h",
+ "typed_allocator.h",
+ "types.h",
+ "variant.h",
+ "variant_encode_decode.h",
+ "variant_op_registry.h",
+ "variant_tensor_data.h",
+ ],
+ visibility = [
+ "//tensorflow/core:__pkg__",
+ "//tensorflow/core/util:__pkg__",
+ ],
+ deps = [
+ ":allocation_description_proto_cc",
+ ":allocator",
+ ":bfloat16",
+ ":log_memory_proto_cc",
+ ":numeric_types",
+ ":resource_handle",
+ ":resource_handle_proto_cc",
+ ":tensor_description_proto_cc",
+ ":tensor_proto_cc",
+ ":tensor_shape",
+ ":tensor_types",
+ ":type_index",
+ ":type_traits",
+ ":types_proto_cc",
+ "//tensorflow/core/lib/core:coding",
+ "//tensorflow/core/lib/core:errors",
+ "//tensorflow/core/lib/core:refcount",
+ "//tensorflow/core/lib/core:status",
+ "//tensorflow/core/lib/core:stringpiece",
+ "//tensorflow/core/lib/gtl:array_slice",
+ "//tensorflow/core/lib/gtl:flatmap",
+ "//tensorflow/core/lib/gtl:inlined_vector",
+ "//tensorflow/core/lib/hash",
+ "//tensorflow/core/lib/strings:str_util",
+ "//tensorflow/core/lib/strings:strcat",
+ "//tensorflow/core/platform:abi",
+ "//tensorflow/core/platform:logging",
+ "//tensorflow/core/platform:macros",
+ "//tensorflow/core/platform:platform_port",
+ "//tensorflow/core/platform:protobuf",
+ "//tensorflow/core/platform:strcat",
+ "//tensorflow/core/platform:tensor_coding",
+ "//tensorflow/core/platform:types",
+ "//tensorflow/core/public:version",
+ "//third_party/eigen3",
+ "@com_google_absl//absl/memory",
+ "@com_google_absl//absl/strings",
+ ],
+ alwayslink = 1,
+)
+
cc_library(
name = "shape_inference",
srcs = ["shape_inference.cc"],
diff --git a/tensorflow/core/framework/tensor_key.h b/tensorflow/core/framework/tensor_key.h
new file mode 100644
index 00000000..8eff58b2
--- /dev/null
+++ b/tensorflow/core/framework/tensor_key.h
@@ -0,0 +1,64 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/tensor.h"
+
+namespace tensorflow {
+
+class TensorKey : public Tensor {
+ public:
+ using Tensor::Tensor;
+
+ TensorKey(const Tensor& t) : Tensor(t) {}
+
+ // Equality operator. Needed for absl hashing.
+ friend bool operator==(const TensorKey& t1, const TensorKey& t2) {
+ if (t1.dtype() != t2.dtype() || t1.shape() != t2.shape()) {
+ return false;
+ }
+ if (DataTypeCanUseMemcpy(t1.dtype())) {
+ return t1.tensor_data() == t2.tensor_data();
+ }
+ if (t1.dtype() == DT_STRING) {
+ const auto s1 = t1.unaligned_flat<tstring>();
+ const auto s2 = t2.unaligned_flat<tstring>();
+ for (int64 i = 0, n = t1.NumElements(); i < n; ++i) {
+ if (TF_PREDICT_FALSE(s1(i) != s2(i))) {
+ return false;
+ }
+ }
+ return true;
+ }
+ return false;
+ }
+
+ friend bool operator!=(const TensorKey& t1, const TensorKey& t2) {
+ return !(t1==t2);
+ }
+
+ // AbslHashValue() function, needed for absl hashing.
+ template <typename H>
+ friend H AbslHashValue(H h, const TensorKey& k) {
+ uint8* d = (uint8*)(k.data());
+ size_t s = k.AllocatedBytes();
+ std::vector<uint8> vec;
+ for (int i=0; i < s; i++) {
+ vec.push_back(d[i]);
+ }
+ return H::combine(std::move(h), s);
+ }
+};
+
+} //namespace tensorflow
\ No newline at end of file
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index f5a480b3..4ef86efb 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3219,6 +3219,7 @@ tf_cc_tests(
],
deps = [
":eigen_helpers",
+ "//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings",
--
2.27.0

70
CVE-2021-29516-2.patch Normal file
View File

@ -0,0 +1,70 @@
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

904
CVE-2021-29516-3.patch Normal file
View File

@ -0,0 +1,904 @@
From be6b1fdb0699d4000b70ad32cc23d1503e5c7511 Mon Sep 17 00:00:00 2001
From: Edward Loper <edloper@google.com>
Date: Wed, 14 Oct 2020 09:41:17 -0700
Subject: [PATCH 1/1] Added gradients for RaggedTensorToVariant and
RaggedTensorFromVariant. (This allows gradients to pass through map_fn when
it is applied to ragged tensors.)
PiperOrigin-RevId: 337108621
Change-Id: I73d5f3296181877f0cc4c7a6273b693bcf8310ab
---
tensorflow/core/kernels/BUILD | 15 ++
.../kernels/ragged_tensor_from_variant_op.cc | 164 +++++++---------
.../kernels/ragged_tensor_to_variant_op.cc | 180 +++++++++++-------
.../core/kernels/ragged_tensor_variant.cc | 86 +++++++++
.../core/kernels/ragged_tensor_variant.h | 110 +++++++++++
tensorflow/core/ops/ragged_conversion_ops.cc | 20 +-
tensorflow/python/ops/ragged/BUILD | 1 +
9 files changed, 478 insertions(+), 172 deletions(-)
create mode 100644 tensorflow/core/framework/tensor_key.h
create mode 100644 tensorflow/core/kernels/ragged_tensor_variant.cc
create mode 100644 tensorflow/core/kernels/ragged_tensor_variant.h
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index f5a480b3..12adb2b2 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1529,10 +1529,22 @@ tf_cc_test(
],
)
+cc_library(
+ name = "ragged_tensor_variant",
+ srcs = ["ragged_tensor_variant.cc"],
+ hdrs = ["ragged_tensor_variant.h"],
+ deps = [
+ ":cwise_op",
+ "//tensorflow/core:framework",
+ ],
+)
+
tf_kernel_library(
name = "ragged_tensor_to_variant_op",
srcs = ["ragged_tensor_to_variant_op.cc"],
deps = [
+ ":concat_lib",
+ ":ragged_tensor_variant",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
],
@@ -1542,6 +1554,7 @@ tf_kernel_library(
name = "ragged_tensor_from_variant_op",
srcs = ["ragged_tensor_from_variant_op.cc"],
deps = [
+ ":ragged_tensor_variant",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
],
@@ -1554,6 +1567,7 @@ tf_cc_test(
deps = [
":ops_testutil",
":ragged_tensor_to_variant_op",
+ ":ragged_tensor_variant",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
@@ -1570,6 +1584,7 @@ tf_cc_test(
deps = [
":ops_testutil",
":ragged_tensor_from_variant_op",
+ ":ragged_tensor_variant",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
diff --git a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
index d7b6a89a..fa8853af 100644
--- a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
@@ -20,110 +20,76 @@ limitations under the License.
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/framework/variant_encode_decode.h"
+#include "tensorflow/core/kernels/ragged_tensor_variant.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace {
-struct RaggedTensor {
- Tensor values;
- std::vector<Tensor> nested_splits;
-};
-
-Status RaggedComponentsFromVariant(const Tensor& encoded_variant,
- int ragged_rank, DataType value_dtype,
- DataType split_dtype,
- std::vector<RaggedTensor>* decoded_ragged) {
+Status RaggedComponentsFromVariant(
+ const Tensor& encoded_variant, int ragged_rank, DataType value_dtype,
+ DataType split_dtype, std::vector<RaggedTensorVariant>* decoded_ragged) {
const auto& flat_variants = encoded_variant.flat<Variant>();
- decoded_ragged->resize(flat_variants.size());
- // Step 1: Extract the 1-D DT_VARIANT Tensor from each Variant element in the
- // input.
+ decoded_ragged->reserve(flat_variants.size());
+
for (int i = 0; i < flat_variants.size(); i++) {
const auto& flat_variant = flat_variants(i);
- const Tensor* encoded_list = flat_variant.get<Tensor>();
- if (encoded_list == nullptr) {
+ const RaggedTensorVariant* decoded =
+ flat_variant.get<RaggedTensorVariant>();
+ if (decoded == nullptr) {
return errors::InvalidArgument(
"Input Variant element at index ", i,
- " doesn't hold a Tensor: ", flat_variant.DebugString());
+ " doesn't hold a RaggedTensorVariant: ", flat_variant.DebugString());
}
- if (encoded_list->dims() != 1) {
+ decoded_ragged->push_back(*decoded);
+ decoded = &decoded_ragged->back();
+ // Check ragged rank & types
+ if (decoded->ragged_rank() != ragged_rank) {
return errors::InvalidArgument(
- "Encoded input Variant must have rank 1, but found rank: ",
- encoded_list->dims(),
- ". encoded input Variant: ", encoded_list->DebugString());
+ "Encoded input RaggedTensorVariant has ragged_rank=",
+ decoded->ragged_rank(), ". Expected ragged_rank=", ragged_rank, ".");
}
- if (encoded_list->NumElements() != (ragged_rank + 1) &&
- encoded_list->NumElements() != 1) {
+ if (decoded->values().dtype() != value_dtype) {
return errors::InvalidArgument(
- "Encoded input Variant must hold either input_ragged_rank + 1 "
- "Tensors or an empty Tensor (zero splits Tensors, 1 values Tensor), "
- "input_ragged_rank: ",
- ragged_rank,
- ", encoded input Variant: ", encoded_list->DebugString());
+ "Expected values Tensor dtype: ", DataTypeString(value_dtype),
+ ", found: ", DataTypeString(decoded->values().dtype()));
}
- const auto& input_vec = encoded_list->vec<Variant>();
-
- // Step 2: Get the splits and value Tensors from the 1-D DT_VARIANT Tensor
- // to create the component RaggedTensors.
- (*decoded_ragged)[i].nested_splits.reserve(ragged_rank);
- for (int j = 0; j < ragged_rank; j++) {
- const Tensor* split_tensor = input_vec(j).get<Tensor>();
- if (split_tensor == nullptr) {
- return errors::InvalidArgument(
- "Encoded scalar element at index ", i,
- " doesn't have a splits Tensor at split_index ", j, ": ",
- input_vec(j).DebugString());
- }
- Tensor splits_tensor = *split_tensor;
- if (splits_tensor.dtype() != split_dtype) {
+ if (decoded->values().dims() < 1) {
+ return errors::InvalidArgument(
+ "Ragged values must have rank >= 1; encoded scalar element at index ",
+ i, " has values Tensor: ", decoded->values().DebugString());
+ }
+ for (const auto& splits : decoded->nested_splits()) {
+ if (splits.dtype() != split_dtype) {
return errors::InvalidArgument(
- "Expected splits Tensor dtype: ", split_dtype,
- ", found: ", splits_tensor.dtype());
+ "Expected row_splits Tensor dtype: ", DataTypeString(split_dtype),
+ ", found: ", DataTypeString(splits.dtype()));
}
- if (splits_tensor.dims() != 1) {
+ if (splits.dims() != 1) {
return errors::InvalidArgument(
"Ragged splits must have rank 1; encoded scalar element at index ",
- i, " has splits Tensor at split_index ", j, ": ",
- splits_tensor.DebugString());
+ i, " has splits Tensor ", splits.DebugString());
}
- (*decoded_ragged)[i].nested_splits.push_back(splits_tensor);
- }
- const Tensor* values_tensor = input_vec(ragged_rank).get<Tensor>();
- if (values_tensor == nullptr) {
- return errors::InvalidArgument("Encoded scalar element at index ", i,
- " doesn't have a values Tensor: ",
- input_vec(ragged_rank).DebugString());
- }
- if (values_tensor->dtype() != value_dtype) {
- return errors::InvalidArgument(
- "Expected values Tensor dtype: ", DataTypeString(value_dtype),
- ", found: ", DataTypeString(values_tensor->dtype()));
- }
- if (values_tensor->dims() < 1) {
- return errors::InvalidArgument(
- "Ragged values must have rank >= 1; encoded scalar element at index ",
- i, " has values Tensor: ", values_tensor->DebugString());
}
- (*decoded_ragged)[i].values = *values_tensor;
}
return Status::OK();
}
template <typename VALUE_TYPE, typename SPLIT_TYPE>
Status NestedStackRaggedTensors(
- const std::vector<RaggedTensor>& ragged_components,
+ const std::vector<RaggedTensorVariant>& ragged_components,
const std::vector<int>& nested_dim_sizes, const int input_ragged_rank,
- const int output_ragged_rank, RaggedTensor* output_ragged) {
- output_ragged->nested_splits.reserve(output_ragged_rank);
+ const int output_ragged_rank, RaggedTensorVariant* output_ragged) {
+ output_ragged->mutable_nested_splits()->reserve(output_ragged_rank);
const int dims = nested_dim_sizes.size();
// Populate first `dims - 1` splits.
for (int i = 0; i < dims - 1; i++) {
int dims_splits_size = nested_dim_sizes[i] + 1;
- output_ragged->nested_splits.push_back(Tensor(
- DataTypeToEnum<SPLIT_TYPE>::value, TensorShape({dims_splits_size})));
- auto splits_vec = output_ragged->nested_splits[i].vec<SPLIT_TYPE>();
+ output_ragged->append_splits(Tensor(DataTypeToEnum<SPLIT_TYPE>::value,
+ TensorShape({dims_splits_size})));
+ auto splits_vec = output_ragged->mutable_splits(i)->vec<SPLIT_TYPE>();
int split_diff = nested_dim_sizes[i + 1];
for (int j = 0; j < dims_splits_size; j++) {
splits_vec(j) = j * split_diff;
@@ -132,15 +98,15 @@ Status NestedStackRaggedTensors(
// Populate `dims`-th split.
int splits_size = ragged_components.size() + 1;
- output_ragged->nested_splits.push_back(
+ output_ragged->append_splits(
Tensor(DataTypeToEnum<SPLIT_TYPE>::value, TensorShape({splits_size})));
auto dims_splits_vec =
- output_ragged->nested_splits[dims - 1].vec<SPLIT_TYPE>();
+ output_ragged->mutable_splits(dims - 1)->vec<SPLIT_TYPE>();
dims_splits_vec(0) = 0;
for (int i = 0; i < ragged_components.size(); i++) {
- int split_val = ragged_components[i].values.shape().dim_size(0);
- if (input_ragged_rank != 0 && !ragged_components[i].nested_splits.empty()) {
- split_val = ragged_components[i].nested_splits[0].NumElements() - 1;
+ int split_val = ragged_components[i].values().shape().dim_size(0);
+ if (input_ragged_rank != 0 && ragged_components[i].ragged_rank() > 0) {
+ split_val = ragged_components[i].splits(0).NumElements() - 1;
}
dims_splits_vec(i + 1) = dims_splits_vec(i) + split_val;
}
@@ -150,24 +116,24 @@ Status NestedStackRaggedTensors(
int split_index = dims + i;
int split_size = 1;
for (int j = 0; j < ragged_components.size(); j++) {
- if (!ragged_components[j].nested_splits.empty()) {
- split_size += ragged_components[j].nested_splits[i].NumElements() - 1;
+ if (!ragged_components[j].nested_splits().empty()) {
+ split_size += ragged_components[j].splits(i).NumElements() - 1;
}
}
- output_ragged->nested_splits.push_back(
+ output_ragged->append_splits(
Tensor(DataTypeToEnum<SPLIT_TYPE>::value, TensorShape({split_size})));
auto splits_vec =
- output_ragged->nested_splits[split_index].vec<SPLIT_TYPE>();
+ output_ragged->mutable_splits(split_index)->vec<SPLIT_TYPE>();
splits_vec(0) = 0;
SPLIT_TYPE last_split_value = 0;
int index = 1;
for (int j = 0; j < ragged_components.size(); j++) {
- if (ragged_components[j].nested_splits.empty()) {
+ if (ragged_components[j].nested_splits().empty()) {
// Corner case: empty row. e.g [ [[x], [x]], [] ]
continue;
}
auto component_splits_vec =
- ragged_components[j].nested_splits[i].vec<SPLIT_TYPE>();
+ ragged_components[j].splits(i).vec<SPLIT_TYPE>();
for (int k = 1; k < component_splits_vec.size(); k++, index++) {
splits_vec(index) = component_splits_vec(k) + last_split_value;
}
@@ -187,35 +153,35 @@ Status NestedStackRaggedTensors(
if (ragged_components.empty()) {
component_values_shape = TensorShape({0});
} else {
- component_values_shape = ragged_components[0].values.shape();
+ component_values_shape = ragged_components[0].values().shape();
}
// Populate values.
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()) {
+ if (ragged_components[i].values().dims() != component_values_shape.dims()) {
return errors::InvalidArgument(
"Rank of values must match for all "
"components; values shape at index 0: ",
component_values_shape.DebugString(), ", values shape at index ", i,
- ": ", ragged_components[i].values.shape().DebugString());
+ ": ", ragged_components[i].values().shape().DebugString());
}
- values_size += ragged_components[i].values.shape().dim_size(0);
+ values_size += ragged_components[i].values().shape().dim_size(0);
}
component_values_shape.set_dim(0, values_size);
- output_ragged->values =
- Tensor(DataTypeToEnum<VALUE_TYPE>::value, component_values_shape);
+ output_ragged->set_values(
+ Tensor(DataTypeToEnum<VALUE_TYPE>::value, component_values_shape));
auto output_values_flat =
- output_ragged->values.flat_outer_dims<VALUE_TYPE, 2>();
+ output_ragged->mutable_values()->flat_outer_dims<VALUE_TYPE, 2>();
int values_index = 0;
for (int i = 0; i < ragged_components.size(); i++) {
auto component_values_flat =
- ragged_components[i].values.flat_outer_dims<VALUE_TYPE, 2>();
- int num_inner_elements = ragged_components[i].values.NumElements();
- if (ragged_components[i].values.dim_size(0) > 0) {
- num_inner_elements /= ragged_components[i].values.dim_size(0);
+ ragged_components[i].values().flat_outer_dims<VALUE_TYPE, 2>();
+ int num_inner_elements = ragged_components[i].values().NumElements();
+ if (ragged_components[i].values().dim_size(0) > 0) {
+ num_inner_elements /= ragged_components[i].values().dim_size(0);
}
- for (int j = 0; j < ragged_components[i].values.dim_size(0);
+ for (int j = 0; j < ragged_components[i].values().dim_size(0);
j++, values_index++) {
for (int k = 0; k < num_inner_elements; k++) {
output_values_flat(values_index, k) = component_values_flat(j, k);
@@ -265,7 +231,7 @@ class RaggedTensorFromVariantOp : public OpKernel {
// Decode all variants.
const auto value_dtype = DataTypeToEnum<VALUE_TYPE>::v();
const auto split_dtype = DataTypeToEnum<SPLIT_TYPE>::v();
- std::vector<RaggedTensor> decoded_components;
+ std::vector<RaggedTensorVariant> decoded_components;
OP_REQUIRES_OK(context, RaggedComponentsFromVariant(
encoded_variant, input_ragged_rank_,
value_dtype, split_dtype, &decoded_components));
@@ -281,7 +247,7 @@ class RaggedTensorFromVariantOp : public OpKernel {
for (int i = 0; i < encoded_variant.dims(); i++) {
encoded_dim_sizes[i] = encoded_variant.dim_size(i);
}
- RaggedTensor output_ragged;
+ RaggedTensorVariant output_ragged;
OP_REQUIRES_OK(
context, NestedStackRaggedTensors<VALUE_TYPE, SPLIT_TYPE>(
decoded_components, encoded_dim_sizes, input_ragged_rank_,
@@ -296,15 +262,15 @@ class RaggedTensorFromVariantOp : public OpKernel {
int output_ragged_rank_;
void ReturnRaggedTensor(OpKernelContext* context,
- RaggedTensor ragged_tensor) {
- int ragged_rank = ragged_tensor.nested_splits.size();
+ const RaggedTensorVariant& ragged_tensor) {
+ int ragged_rank = ragged_tensor.ragged_rank();
OpOutputList splits_out;
OP_REQUIRES_OK(context,
context->output_list("output_nested_splits", &splits_out));
for (int i = 0; i < ragged_rank; i++) {
- splits_out.set(i, ragged_tensor.nested_splits[i]);
+ splits_out.set(i, ragged_tensor.splits(i));
}
- context->set_output(ragged_rank, ragged_tensor.values);
+ context->set_output(ragged_rank, ragged_tensor.values());
}
};
diff --git a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
index 3190534b..a60e5c62 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
@@ -18,50 +18,38 @@ limitations under the License.
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/framework/variant_encode_decode.h"
+#include "tensorflow/core/framework/variant_op_registry.h"
+#include "tensorflow/core/kernels/concat_lib.h"
+#include "tensorflow/core/kernels/ragged_tensor_variant.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/util/tensor_ops_util.h"
namespace tensorflow {
namespace {
-struct RaggedTensor {
- Tensor values;
- std::vector<Tensor> nested_splits;
-};
-
-Status RaggedToVariant(const RaggedTensor& ragged, Tensor* encoded_list) {
- // Encode as a rank-1 Variant Tensor.
- int ragged_rank = ragged.nested_splits.size();
- *encoded_list = Tensor(DT_VARIANT, TensorShape({ragged_rank + 1}));
- auto encoded_vec = encoded_list->vec<Variant>();
- for (int i = 0; i < ragged_rank; i++) {
- encoded_vec(i) = ragged.nested_splits[i];
- }
- encoded_vec(ragged_rank) = ragged.values;
- return Status::OK();
-}
-
template <typename VALUE_TYPE, typename SPLIT_TYPE>
-Status UnbatchRaggedZerothDim(const RaggedTensor& batched_ragged,
- std::vector<RaggedTensor>* ragged_components) {
+Status UnbatchRaggedZerothDim(
+ const RaggedTensorVariant& batched_ragged,
+ std::vector<RaggedTensorVariant>* ragged_components) {
// Set up the component Ragged Tensors.
- int ragged_rank = batched_ragged.nested_splits.size();
- auto batched_splits_top_vec =
- batched_ragged.nested_splits[0].vec<SPLIT_TYPE>();
+ int ragged_rank = batched_ragged.ragged_rank();
+ auto batched_splits_top_vec = batched_ragged.splits(0).vec<SPLIT_TYPE>();
int num_components = batched_splits_top_vec.size() - 1;
int num_splits = ragged_rank - 1;
ragged_components->resize(num_components);
- for (RaggedTensor ragged_component : *ragged_components) {
- ragged_component.nested_splits.reserve(num_splits);
+ for (RaggedTensorVariant& ragged_component : *ragged_components) {
+ ragged_component.mutable_nested_splits()->reserve(num_splits);
}
- const auto& batched_flat = batched_ragged.values.flat<VALUE_TYPE>();
- int num_inner_elems = batched_ragged.values.NumElements();
- if (batched_ragged.values.dim_size(0) > 1) {
- num_inner_elems /= batched_ragged.values.dim_size(0);
+ const auto& batched_flat = batched_ragged.values().flat<VALUE_TYPE>();
+ int num_inner_elems = batched_ragged.values().NumElements();
+ if (batched_ragged.values().dim_size(0) > 1) {
+ num_inner_elems /= batched_ragged.values().dim_size(0);
}
- TensorShape values_shape = batched_ragged.values.shape();
+ TensorShape values_shape = batched_ragged.values().shape();
// Corner case: ragged_rank == 1, e.g. [[1, 2, 3], [4, 5]]
if (num_splits == 0) {
@@ -70,10 +58,10 @@ Status UnbatchRaggedZerothDim(const RaggedTensor& batched_ragged,
int limit = batched_splits_top_vec(i + 1);
int num_values = limit - start;
values_shape.set_dim(0, num_values);
- (*ragged_components)[i].values =
- Tensor(DataTypeToEnum<VALUE_TYPE>::value, values_shape);
+ (*ragged_components)[i].set_values(
+ Tensor(DataTypeToEnum<VALUE_TYPE>::value, values_shape));
auto ragged_component_values_flat =
- (*ragged_components)[i].values.flat<VALUE_TYPE>();
+ (*ragged_components)[i].mutable_values()->flat<VALUE_TYPE>();
for (int j = 0; j < num_values * num_inner_elems; j++) {
ragged_component_values_flat(j) =
batched_flat(j + start * num_inner_elems);
@@ -86,8 +74,7 @@ Status UnbatchRaggedZerothDim(const RaggedTensor& batched_ragged,
std::vector<typename TTypes<SPLIT_TYPE>::ConstVec> batched_splits_vec;
batched_splits_vec.reserve(ragged_rank);
for (int i = 0; i < ragged_rank; i++) {
- batched_splits_vec.push_back(
- batched_ragged.nested_splits[i].vec<SPLIT_TYPE>());
+ batched_splits_vec.push_back(batched_ragged.splits(i).vec<SPLIT_TYPE>());
}
std::vector<int> index(num_splits, 1);
std::vector<int> ragged_component_values_size(num_components, 0);
@@ -104,10 +91,10 @@ Status UnbatchRaggedZerothDim(const RaggedTensor& batched_ragged,
int last_index = ragged_component_splits_vec[j - 1].size() - 1;
split_size = ragged_component_splits_vec[j - 1](last_index) + 1;
}
- (*ragged_components)[i].nested_splits.push_back(
+ (*ragged_components)[i].append_splits(
Tensor(DataTypeToEnum<SPLIT_TYPE>::value, TensorShape({split_size})));
ragged_component_splits_vec.push_back(
- (*ragged_components)[i].nested_splits[j].vec<SPLIT_TYPE>());
+ (*ragged_components)[i].mutable_splits(j)->vec<SPLIT_TYPE>());
SPLIT_TYPE last_split_value = batched_splits_vec[j + 1](index[j] - 1);
ragged_component_splits_vec[j](0) = 0;
for (int k = 1; k < split_size; k++, index[j]++) {
@@ -125,10 +112,10 @@ Status UnbatchRaggedZerothDim(const RaggedTensor& batched_ragged,
for (int i = 0; i < num_components; i++) {
int num_values = ragged_component_values_size[i];
values_shape.set_dim(0, num_values);
- (*ragged_components)[i].values =
- Tensor(DataTypeToEnum<VALUE_TYPE>::value, values_shape);
+ (*ragged_components)[i].set_values(
+ Tensor(DataTypeToEnum<VALUE_TYPE>::value, values_shape));
auto ragged_component_values_flat =
- (*ragged_components)[i].values.flat<VALUE_TYPE>();
+ (*ragged_components)[i].mutable_values()->flat<VALUE_TYPE>();
for (int j = 0; j < num_values * num_inner_elems; j++, value_index++) {
ragged_component_values_flat(j) = batched_flat(value_index);
}
@@ -152,24 +139,21 @@ class RaggedTensorToVariantOp : public OpKernel {
OP_REQUIRES_OK(context, context->input_list("rt_nested_splits",
&ragged_nested_splits_in));
const int ragged_nested_splits_len = ragged_nested_splits_in.size();
- RaggedTensor batched_ragged_input;
+ RaggedTensorVariant batched_ragged_input;
// Read ragged_values input.
- batched_ragged_input.values = context->input(ragged_nested_splits_len);
- batched_ragged_input.nested_splits.reserve(ragged_nested_splits_len);
+ batched_ragged_input.set_values(context->input(ragged_nested_splits_len));
+ batched_ragged_input.mutable_nested_splits()->reserve(
+ ragged_nested_splits_len);
for (int i = 0; i < ragged_nested_splits_len; i++) {
- batched_ragged_input.nested_splits.push_back(ragged_nested_splits_in[i]);
+ batched_ragged_input.append_splits(ragged_nested_splits_in[i]);
}
if (!batched_input_) {
- // Encode the input as is.
- Tensor encoded_list;
- OP_REQUIRES_OK(context,
- RaggedToVariant(batched_ragged_input, &encoded_list));
// Encode as a Scalar Variant Tensor.
Tensor* encoded_scalar;
OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape({}),
&encoded_scalar));
- encoded_scalar->scalar<Variant>()() = std::move(encoded_list);
+ encoded_scalar->scalar<Variant>()() = std::move(batched_ragged_input);
return;
}
@@ -180,24 +164,19 @@ class RaggedTensorToVariantOp : public OpKernel {
"received rt_nested_splits of length 0."));
// Unbatch the Ragged Tensor and encode the components.
- std::vector<RaggedTensor> ragged_components;
+ std::vector<RaggedTensorVariant> unbatched_ragged_input;
OP_REQUIRES_OK(context, UnbatchRaggedZerothDim<VALUE_TYPE, SPLIT_TYPE>(
- batched_ragged_input, &ragged_components));
- std::vector<Tensor> encoded_components(ragged_components.size());
- for (int i = 0; i < ragged_components.size(); i++) {
- OP_REQUIRES_OK(context, RaggedToVariant(ragged_components[i],
- &encoded_components[i]));
- }
+ batched_ragged_input, &unbatched_ragged_input));
// Bundle the encoded scalar Variant Tensors into a rank-1 Variant Tensor.
- Tensor* encoded_ragged;
- int output_size = ragged_components.size();
+ Tensor* encoded_vector;
+ int output_size = unbatched_ragged_input.size();
OP_REQUIRES_OK(context,
context->allocate_output(0, TensorShape({output_size}),
- &encoded_ragged));
- auto encoded_ragged_vec = encoded_ragged->vec<Variant>();
+ &encoded_vector));
+ auto encoded_vector_t = encoded_vector->vec<Variant>();
for (int i = 0; i < output_size; i++) {
- encoded_ragged_vec(i) = encoded_components[i];
+ encoded_vector_t(i) = unbatched_ragged_input[i];
}
}
@@ -205,12 +184,81 @@ class RaggedTensorToVariantOp : public OpKernel {
bool batched_input_;
};
-#define REGISTER_KERNELS_WITH_SPLIT_TYPE(value_type, split_type) \
- REGISTER_KERNEL_BUILDER(Name("RaggedTensorToVariant") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<value_type>("Tvalues") \
- .TypeConstraint<split_type>("Tsplits"), \
- RaggedTensorToVariantOp<value_type, split_type>);
+template <typename VALUE_TYPE, typename SPLIT_TYPE>
+class RaggedTensorToVariantGradientOp : public OpKernel {
+ public:
+ using OpKernel::OpKernel;
+
+ void Compute(OpKernelContext* context) override {
+ // Read inputs.
+ Tensor encoded_variant = context->input(0);
+ Tensor row_splits = context->input(1);
+ auto flat_row_splits = row_splits.flat<SPLIT_TYPE>();
+ TensorShape dense_values_shape;
+ OP_REQUIRES_OK(context,
+ TensorShapeUtils::MakeShape(context->input(2).vec<int32>(),
+ &dense_values_shape));
+
+ const auto& flat_variants = encoded_variant.flat<Variant>();
+
+ // Get a Tensor containing the flat_values for each variant.
+ std::vector<Tensor> values;
+ for (int i = 0; i < flat_variants.size(); ++i) {
+ if (const auto* encoded = flat_variants(i).get<RaggedTensorVariant>()) {
+ values.push_back(encoded->values());
+ } else {
+ // Missing value: this happens if only some of the variant values
+ // generated by ragged_tensor_to_variant impacted the value that we're
+ // calculating the gradient for. In this case, we will see a
+ // default-constructed variant; so treat it as a zero tensor with the
+ // appropriate shape.
+ const auto value_dtype = DataTypeToEnum<VALUE_TYPE>::v();
+ int piece_size = flat_row_splits(i + 1) - flat_row_splits(i);
+ TensorShape zeros_shape = dense_values_shape;
+ zeros_shape.set_dim(0, piece_size);
+ Tensor zero(value_dtype, zeros_shape);
+ zero.flat<VALUE_TYPE>() =
+ zero.flat<VALUE_TYPE>().constant(VALUE_TYPE());
+ values.push_back(zero);
+ }
+ }
+
+ if (values.size() == 1) {
+ // Just one flat_value tensor: return as-is.
+ context->set_output(0, values[0]);
+ } else {
+ // Multiple flat_values tensors: concatenate them together.
+ using Piece = typename TTypes<VALUE_TYPE, 2>::Matrix;
+ using ConstPiece = typename TTypes<VALUE_TYPE, 2>::ConstMatrix;
+ std::vector<std::unique_ptr<ConstPiece>> pieces;
+ pieces.reserve(values.size());
+ for (const Tensor& t : values) {
+ pieces.emplace_back(
+ new ConstPiece(t.shaped<VALUE_TYPE, 2>({1, t.NumElements()})));
+ }
+ Tensor* out = nullptr;
+ OP_REQUIRES_OK(context,
+ context->allocate_output(0, dense_values_shape, &out));
+ Piece out_flat =
+ out->shaped<VALUE_TYPE, 2>({1, dense_values_shape.num_elements()});
+ ConcatCPU<VALUE_TYPE>(context->device(), pieces, &out_flat);
+ }
+ }
+};
+
+#define REGISTER_KERNELS_WITH_SPLIT_TYPE(value_type, split_type) \
+ REGISTER_KERNEL_BUILDER(Name("RaggedTensorToVariant") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<value_type>("Tvalues") \
+ .TypeConstraint<split_type>("Tsplits"), \
+ RaggedTensorToVariantOp<value_type, split_type>); \
+ REGISTER_KERNEL_BUILDER( \
+ Name("RaggedTensorToVariantGradient") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<value_type>("Tvalues") \
+ .TypeConstraint<split_type>("Tsplits"), \
+ RaggedTensorToVariantGradientOp<value_type, split_type>);
+
#define REGISTER_KERNELS(value_type) \
REGISTER_KERNELS_WITH_SPLIT_TYPE(value_type, int32) \
REGISTER_KERNELS_WITH_SPLIT_TYPE(value_type, int64)
diff --git a/tensorflow/core/kernels/ragged_tensor_variant.cc b/tensorflow/core/kernels/ragged_tensor_variant.cc
new file mode 100644
index 00000000..94663138
--- /dev/null
+++ b/tensorflow/core/kernels/ragged_tensor_variant.cc
@@ -0,0 +1,86 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#define EIGEN_USE_THREADS
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#define EIGEN_USE_GPU
+#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include "tensorflow/core/kernels/ragged_tensor_variant.h"
+
+namespace tensorflow {
+
+string RaggedTensorVariant::TypeName() const { return "RaggedTensorVariant"; }
+
+string RaggedTensorVariant::DebugString() const {
+ return absl::StrCat(
+ "RaggedTensorVariant(dtype=", DataTypeString(values_.dtype()),
+ ", ragged_rank=", nested_splits_.size(), ", splits_dtype=",
+ DataTypeString(nested_splits_.empty() ? DT_INVALID
+ : nested_splits_.back().dtype()));
+}
+
+void RaggedTensorVariant::Encode(VariantTensorData* data) const {
+ data->set_type_name(TypeName());
+ for (const auto& splits : nested_splits_) {
+ *data->add_tensors() = splits;
+ }
+ *data->add_tensors() = values_;
+}
+
+bool RaggedTensorVariant::Decode(const VariantTensorData& data) {
+ if (data.tensors_size() < 1) {
+ return false;
+ }
+ nested_splits_.assign(data.tensors().begin(),
+ std::prev(data.tensors().end()));
+ values_ = data.tensors().back();
+ return true;
+}
+
+namespace {
+
+Status RaggedTensorVariantDeviceCopy(
+ const RaggedTensorVariant& from, RaggedTensorVariant* to,
+ const UnaryVariantOpRegistry::AsyncTensorDeviceCopyFn& copy) {
+ TF_RETURN_IF_ERROR(copy(from.values(), to->mutable_values()));
+ // TODO(b/170415165) Should we use `copy` to move splits from device<->host?
+ *to->mutable_nested_splits() = from.nested_splits();
+ return Status::OK();
+}
+
+} // namespace
+
+REGISTER_UNARY_VARIANT_UNARY_OP_FUNCTION(
+ ZEROS_LIKE_VARIANT_UNARY_OP, DEVICE_CPU, RaggedTensorVariant,
+ RaggedTensorVariantZerosLike<CPUDevice>);
+
+REGISTER_UNARY_VARIANT_BINARY_OP_FUNCTION(
+ ADD_VARIANT_BINARY_OP, DEVICE_CPU, RaggedTensorVariant,
+ RaggedTensorVariantBinaryAdd<CPUDevice>);
+
+REGISTER_UNARY_VARIANT_DECODE_FUNCTION(RaggedTensorVariant,
+ "RaggedTensorVariant");
+
+#define REGISTER_RAGGED_TENSOR_VARIANT_COPY(DIRECTION) \
+ INTERNAL_REGISTER_UNARY_VARIANT_DEVICE_COPY_FUNCTION( \
+ RaggedTensorVariant, DIRECTION, RaggedTensorVariantDeviceCopy)
+
+REGISTER_RAGGED_TENSOR_VARIANT_COPY(VariantDeviceCopyDirection::HOST_TO_DEVICE);
+REGISTER_RAGGED_TENSOR_VARIANT_COPY(VariantDeviceCopyDirection::DEVICE_TO_HOST);
+REGISTER_RAGGED_TENSOR_VARIANT_COPY(
+ VariantDeviceCopyDirection::DEVICE_TO_DEVICE);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/ragged_tensor_variant.h b/tensorflow/core/kernels/ragged_tensor_variant.h
new file mode 100644
index 00000000..730758a3
--- /dev/null
+++ b/tensorflow/core/kernels/ragged_tensor_variant.h
@@ -0,0 +1,110 @@
+#include "tensorflow/core/framework/tensor_key.h"
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_KERNELS_RAGGED_TENSOR_VARIANT_H_
+#define TENSORFLOW_CORE_KERNELS_RAGGED_TENSOR_VARIANT_H_
+
+#define EIGEN_USE_THREADS
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#define EIGEN_USE_GPU
+#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include <vector>
+
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/variant_op_registry.h"
+#include "tensorflow/core/framework/variant_tensor_data.h"
+#include "tensorflow/core/kernels/cwise_ops_common.h"
+#include "tensorflow/core/util/tensor_ops_util.h"
+
+namespace tensorflow {
+
+// Class used to store a RaggedTensor as a Variant scalar.
+class RaggedTensorVariant {
+ public:
+ RaggedTensorVariant() {}
+ RaggedTensorVariant(Tensor values, const std::vector<Tensor>& nested_splits)
+ : values_(std::move(values)), nested_splits_(nested_splits) {}
+
+ // Variant support methods.
+ string TypeName() const;
+ string DebugString() const;
+ void Encode(VariantTensorData* data) const;
+ bool Decode(const VariantTensorData& data);
+
+ // The flat_values of the RaggedTensor.
+ const Tensor& values() const { return values_; }
+ Tensor* mutable_values() { return &values_; }
+ void set_values(const Tensor& new_values) { values_ = new_values; }
+
+ // The nested row_splits of the RaggedTensor.
+ int ragged_rank() const { return nested_splits_.size(); }
+ const std::vector<Tensor>& nested_splits() const { return nested_splits_; }
+ std::vector<Tensor>* mutable_nested_splits() { return &nested_splits_; }
+ const Tensor& splits(int i) const { return nested_splits_[i]; }
+ Tensor* mutable_splits(int i) { return &nested_splits_[i]; }
+ void set_nested_splits(const std::vector<Tensor>& nested_splits) {
+ nested_splits_ = nested_splits;
+ }
+ void append_splits(const Tensor& splits) { nested_splits_.push_back(splits); }
+
+ private:
+ Tensor values_;
+ std::vector<Tensor> nested_splits_;
+};
+
+template <typename Device>
+Status RaggedTensorVariantZerosLike(OpKernelContext* c,
+ const RaggedTensorVariant& x,
+ RaggedTensorVariant* y) {
+ y->set_nested_splits(x.nested_splits());
+ TF_RETURN_IF_ERROR(
+ ZerosLikeTensor<Device>(c, x.values(), y->mutable_values()));
+ return Status::OK();
+}
+
+template <typename Device>
+Status RaggedTensorVariantBinaryAdd(OpKernelContext* c,
+ const RaggedTensorVariant& x,
+ const RaggedTensorVariant& y,
+ RaggedTensorVariant* out) {
+ if (x.values().dtype() != y.values().dtype()) {
+ return errors::InvalidArgument(
+ "Can't add RaggedTensorVariants of different dtypes. One is ",
+ DataTypeString(x.values().dtype()), " and the other is ",
+ DataTypeString(y.values().dtype()));
+ }
+ if (x.ragged_rank() != y.ragged_rank()) {
+ return errors::InvalidArgument(
+ "Can't add RaggedTensorVariants of different ragged rank. ", "One is ",
+ x.ragged_rank(), " and the other is ", y.ragged_rank());
+ }
+ for (int i = 0; i < x.ragged_rank(); ++i) {
+ if (TensorKey(x.splits(i)) != TensorKey(y.splits(i))) {
+ return errors::InvalidArgument(
+ "Can't add RaggedTensorVariants with different row_splits.");
+ }
+ }
+ out->set_nested_splits(x.nested_splits());
+ TF_RETURN_IF_ERROR(BinaryAddTensors<Device>(c, x.values(), y.values(),
+ out->mutable_values()));
+ return Status::OK();
+}
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CORE_KERNELS_RAGGED_TENSOR_VARIANT_H_
diff --git a/tensorflow/core/ops/ragged_conversion_ops.cc b/tensorflow/core/ops/ragged_conversion_ops.cc
index 6bee189c..8512bcf3 100644
--- a/tensorflow/core/ops/ragged_conversion_ops.cc
+++ b/tensorflow/core/ops/ragged_conversion_ops.cc
@@ -92,7 +92,8 @@ tensorflow::Status ValidateRowPartitionTypesAndShapes(
Status RaggedTensorToSparseShapeFn(InferenceContext* c);
Status RaggedTensorToVariantShapeFn(InferenceContext* c);
Status RaggedTensorFromVariantShapeFn(InferenceContext* c);
-tensorflow::Status RaggedTensorToTensorShapeFn(InferenceContext* c);
+Status RaggedTensorToVariantGradientShapeFn(InferenceContext* c);
+Status RaggedTensorToTensorShapeFn(InferenceContext* c);
//==============================================================================
// Registered Ops
@@ -129,6 +130,15 @@ REGISTER_OP("RaggedTensorFromVariant")
.Attr("Tsplits: {int32, int64} = DT_INT64")
.SetShapeFn(RaggedTensorFromVariantShapeFn);
+REGISTER_OP("RaggedTensorToVariantGradient")
+ .Input("encoded_ragged_grad: variant")
+ .Input("row_splits: Tsplits")
+ .Input("dense_values_shape: int32")
+ .Output("dense_values_grad: Tvalues")
+ .Attr("Tvalues: type")
+ .Attr("Tsplits: {int32, int64} = DT_INT64")
+ .SetShapeFn(RaggedTensorToVariantGradientShapeFn);
+
REGISTER_OP("RaggedTensorToTensor")
.Attr("T: type")
.Attr("Tindex: {int64, int32}")
@@ -201,6 +211,14 @@ Status RaggedTensorToVariantShapeFn(InferenceContext* c) {
return Status::OK();
}
+Status RaggedTensorToVariantGradientShapeFn(InferenceContext* c) {
+ ShapeHandle shape;
+ TF_RETURN_IF_ERROR(
+ c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(2, &shape));
+ c->set_output(0, shape);
+ return Status::OK();
+}
+
Status RaggedTensorFromVariantShapeFn(InferenceContext* c) {
int64 input_ragged_rank;
TF_RETURN_IF_ERROR(
diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index 95e5602a..34372160 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -507,6 +507,7 @@ py_test(
"//tensorflow/python:framework_ops",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python:tensor_array_grad",
"//tensorflow/python:tensor_shape",
"//tensorflow/python:tensor_spec",
"//tensorflow/python/data/ops:dataset_ops",
--
2.27.0

30
CVE-2021-29516-4.patch Normal file
View File

@ -0,0 +1,30 @@
From b055b9c474cd376259dde8779908f9eeaf097d93 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Tue, 13 Apr 2021 14:49:50 -0700
Subject: [PATCH] Fix `tf.raw_ops.RaggedTensorToVariant` invalid resize.
PiperOrigin-RevId: 368299574
Change-Id: I751c186325aa0bab397928845e790e60c2d90918
---
tensorflow/core/kernels/ragged_tensor_to_variant_op.cc | 5 +++++
1 file changed, 5 insertions(+)
diff --git a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
index a60e5c62..fb1f25fc 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
@@ -165,6 +165,11 @@ class RaggedTensorToVariantOp : public OpKernel {
// Unbatch the Ragged Tensor and encode the components.
std::vector<RaggedTensorVariant> unbatched_ragged_input;
+ auto batched_splits_top_vec =
+ batched_ragged_input.splits(0).vec<SPLIT_TYPE>();
+ int num_components = batched_splits_top_vec.size() - 1;
+ OP_REQUIRES(context, num_components >= 0,
+ errors::Internal("Invalid split argument."));
OP_REQUIRES_OK(context, UnbatchRaggedZerothDim<VALUE_TYPE, SPLIT_TYPE>(
batched_ragged_input, &unbatched_ragged_input));
--
2.27.0

88
CVE-2021-37679.patch Normal file
View File

@ -0,0 +1,88 @@
From 4e2565483d0ffcadc719bd44893fb7f609bb5f12 Mon Sep 17 00:00:00 2001
From: Edward Loper <edloper@google.com>
Date: Thu, 29 Jul 2021 09:50:01 -0700
Subject: [PATCH] Fix bug that could cause map_fn to produce incorrect results
(rather than an error) when mapping over a ragged tensor with an
inappropriate fn_output_signature. (Note: there are cases where the default
value for fn_output_signature is not appropriate, so the user needs to
explicitly specify the correct output signature.)
PiperOrigin-RevId: 387606546
Change-Id: Ib4ea27b9634e6ab413f211cfe809a69a90f0e2cd
---
.../kernels/ragged_tensor_from_variant_op.cc | 16 +++++++++++++
.../ops/ragged/ragged_map_fn_op_test.py | 23 +++++++++++++++++++
2 files changed, 39 insertions(+)
diff --git a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
index d9993bb6d3907..c481d90638e4e 100644
--- a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
@@ -174,7 +174,23 @@ Status NestedStackRaggedTensors(
auto output_values_flat =
output_ragged->mutable_values()->flat_outer_dims<VALUE_TYPE, 2>();
int values_index = 0;
+
+ TensorShape expected_value_shape = component_values_shape;
+ expected_value_shape.RemoveDim(0);
+
for (int i = 0; i < ragged_components.size(); i++) {
+ // Check that the flat_values tensor shape is compatible.
+ TensorShape value_shape = ragged_components[i].values().shape();
+ value_shape.RemoveDim(0);
+ if (value_shape != expected_value_shape) {
+ return errors::InvalidArgument(
+ "All flat_values must have compatible shapes. Shape at index 0: ",
+ expected_value_shape, ". Shape at index ", i, ": ", value_shape,
+ ". If you are using tf.map_fn, then you may need to specify an "
+ "explicit fn_output_signature with appropriate ragged_rank, and/or "
+ "convert output tensors to RaggedTensors.");
+ }
+
auto component_values_flat =
ragged_components[i].values().flat_outer_dims<VALUE_TYPE, 2>();
int num_inner_elements = ragged_components[i].values().NumElements();
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 bead4923a0a4c..ace724ac8711d 100644
--- a/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py
@@ -21,9 +21,11 @@
import numpy as np
from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import map_fn as map_fn_lib
from tensorflow.python.ops import math_ops as mo
from tensorflow.python.ops import string_ops
from tensorflow.python.ops.ragged import ragged_factory_ops
@@ -309,6 +311,27 @@ def testMapOnSparseTensor(self):
)
self.assertAllEqual(id_t2, [[0, 5], [0, 4]])
+ def testRaggedMapWithIncorrectFnOutputSignature(self):
+ x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]])
+ with self.assertRaisesRegex(errors.InvalidArgumentError,
+ 'All flat_values must have compatible shapes'):
+ y = map_fn_lib.map_fn(lambda r: map_fn_lib.map_fn(lambda y: r, r), x)
+ self.evaluate(y)
+
+ def testNestedRaggedMapWithFnOutputSignature(self):
+ ragged1d = ragged_tensor.RaggedTensorSpec([None], dtypes.int32)
+ ragged2d = ragged_tensor.RaggedTensorSpec([None, None], dtypes.int32)
+
+ x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]])
+ # pylint: disable=g-long-lambda
+ y = map_fn_lib.map_fn(
+ lambda r: map_fn_lib.map_fn(
+ lambda y: r, r, fn_output_signature=ragged1d),
+ x,
+ fn_output_signature=ragged2d)
+ expected = [[[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], [[1]]]
+ self.assertAllEqual(y, expected)
+
if __name__ == '__main__':
googletest.main()

10
tensorflow.spec
View File

@ -1,7 +1,7 @@
%global _empty_manifest_terminate_build 0
Name: tensorflow
Version: 2.3.1
Release: 10
Release: 11
Summary: An Open Source Machine Learning Framework for Everyone
License: Apache License 2.0
URL: https://www.tensorflow.org/
@ -180,6 +180,11 @@ Patch0168: CVE-2021-29551.patch
Patch0169: CVE-2021-37645.patch
Patch0170: CVE-2021-37681-1.patch
Patch0171: CVE-2021-37681-2.patch
Patch0172: CVE-2021-29516-1.patch
Patch0173: CVE-2021-29516-2.patch
Patch0174: CVE-2021-29516-3.patch
Patch0175: CVE-2021-29516-4.patch
Patch0176: CVE-2021-37679.patch
Requires: python3-future
Requires: python3-numpy
@ -226,6 +231,9 @@ bazel --output_user_root=`pwd`/../output_user_root build --host_copt=-Wno-string
%{_bindir}/*
%changelog
* Mon Sep 13 2021 yaoxin <yaoxin30@huawei.com> - 2.3.1-11
- Fix CVE-2021-29516 CVE-2021-37679
* Fri Sep 10 2021 yaoxin <yaoxin30@huawei.com> - 2.3.1-10
- Fix CVE-2020-26267 CVE-2021-29515 CVE-2021-29551 CVE-2021-37645 CVE-2021-37681