Make `MakeShape` and `VectorToShape` robust to negative indices.

This avoids a CHECK failure if invalid input data is passed to some
ops. It also tightens up the Python class to reject negative dimensions.
Change: 110401558

1 files changed, 8 insertions, 4 deletions

diff --git a/tensorflow/core/kernels/random_op.cc b/tensorflow/core/kernels/random_op.cc
index c0f8f77d1d..bb1566b4e8 100644
--- a/tensorflow/core/kernels/random_op.cc
+++ b/tensorflow/core/kernels/random_op.cc
@@ -187,12 +187,16 @@ static Status AllocateOutputWithShape(OpKernelContext* ctx, const Tensor& shape,
   }
   if (shape.dtype() == DataType::DT_INT32) {
     auto vec = shape.flat<int32>();
-    TF_RETURN_IF_ERROR(ctx->allocate_output(
-        index, TensorShapeUtils::MakeShape(vec.data(), vec.size()), output));
+    TensorShape tensor_shape;
+    TF_RETURN_IF_ERROR(
+        TensorShapeUtils::MakeShape(vec.data(), vec.size(), &tensor_shape));
+    TF_RETURN_IF_ERROR(ctx->allocate_output(index, tensor_shape, output));
   } else if (shape.dtype() == DataType::DT_INT64) {
     auto vec = shape.flat<int64>();
-    TF_RETURN_IF_ERROR(ctx->allocate_output(
-        index, TensorShapeUtils::MakeShape(vec.data(), vec.size()), output));
+    TensorShape tensor_shape;
+    TF_RETURN_IF_ERROR(
+        TensorShapeUtils::MakeShape(vec.data(), vec.size(), &tensor_shape));
+    TF_RETURN_IF_ERROR(ctx->allocate_output(index, tensor_shape, output));
   } else {
     return errors::InvalidArgument("shape must be a vector of {int32,int64}.");
   }