Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix abort caused by allocating a too large vector.
We need to make sure that the number of dimensions in a shape is within limits.

PiperOrigin-RevId: 408997911
Change-Id: If59e1c23f2ec9c2d4ff4d8632fd62b2a7773a4eb
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Nov 10, 2021
1 parent f6e7c84 commit 1361fb7
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions tensorflow/core/framework/shape_inference.cc
Expand Up @@ -14,6 +14,8 @@ limitations under the License.
==============================================================================*/
#include "tensorflow/core/framework/shape_inference.h"

#include <cstdint>

#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/full_type_util.h"
#include "tensorflow/core/framework/node_def.pb.h"
Expand Down Expand Up @@ -789,6 +791,19 @@ Status InferenceContext::InternalMakeShapeFromTensor(
return ReturnUnknownShape(out);
}
const auto num_dims = Value(shape_dim);
// TODO(mihaimaruseac): Should be `TensorShape::MaxDimensions()` as we are
// not able to materialize shapes with more than this number of dimensions
// but then shape inference would fail for operations such as
// `tf.range`/`tf.ones`, etc. where the shape is not really materialized,
// only used during the inference. Hence, just prevent doing a `reserve`
// with a very large argument.
const int64_t max_dimensions = 1 << 20;
if (num_dims >= max_dimensions) {
return errors::Internal(
"Cannot create a tensor with ", num_dims,
" dimensions, as these would be more than maximum of ",
max_dimensions);
}
std::vector<DimensionHandle> dims;
dims.reserve(num_dims);
for (int i = 0; i < num_dims; i++) dims.push_back(UnknownDim());
Expand Down

0 comments on commit 1361fb7

Please sign in to comment.