OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2026-01-15 12:15:41 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Merge pull request #50773 from tensorflow/mm-cherrypick-fix-onr2.4

Browse Source 
Fix segmentation fault in shape inference logic.
This commit is contained in:
Mihai Maruseac
2021-07-14 12:58:20 -07:00
committed by GitHub
parent 8317a9efdb 9eb9112aa6
commit 2474fca792
1 changed files with 19 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
tensorflow/core/common_runtime/shape_refiner.cc
View File

@@ -120,9 +120,26 @@ Status ShapeRefiner::InferShapesForFunctionSubNode(
TF_RETURN_IF_ERROR(outer_context->MakeShapeFromShapeProto(proto, &handle));
outer_context->set_output(index, handle);
auto* resource = node_context->input_handle_shapes_and_types(0);
const std::vector<ShapeAndType>* resource =
node_context->input_handle_shapes_and_types(0);
if (resource) {
outer_context->set_output_handle_shapes_and_types(index, *resource);
// `ShapesAndType`s contain `ShapeHandle`s. These `ShapeHandle`s point
// to `Shape`s that are owned by a different inference context too. We
// need to copy them to the outer context to prevent them from being
// destroyed before they are used.
std::vector<ShapeAndType> copied_shapes_and_types;
for (auto& shape_and_type : *resource) {
ShapeHandle handle;
TensorShapeProto proto;
node_context->ShapeHandleToProto(shape_and_type.shape, &proto);
TF_RETURN_IF_ERROR(
outer_context->MakeShapeFromShapeProto(proto, &handle));
copied_shapes_and_types.push_back(
ShapeAndType(handle, shape_and_type.dtype, shape_and_type.type));
}
outer_context->set_output_handle_shapes_and_types(
index, copied_shapes_and_types);
}
}