Abort caused by allocating a vector that is too large in Tensorflow

@@ -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"
@@ -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());