27 lines
1.2 KiB
Diff
27 lines
1.2 KiB
Diff
|
From 8a6e874437670045e6c7dc6154c7412b4a2135e2 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Laura Pak <lpak@google.com>
|
||
|
|
Date: Fri, 9 Jul 2021 17:32:55 -0700
|
||
|
|
Subject: [PATCH] Validate num_elements input in tf.raw_ops.TensorListReserve
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 383954564
|
||
|
|
Change-Id: I454bd78eff85bc4f16ddb7e608596971cca47f8f
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/list_kernels.cc | 4 ++++
|
||
|
|
1 file changed, 4 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/list_kernels.cc b/tensorflow/core/kernels/list_kernels.cc
|
||
|
|
index 9a2f373f5ce0c..488e02337f707 100644
|
||
|
|
--- a/tensorflow/core/kernels/list_kernels.cc
|
||
|
|
+++ b/tensorflow/core/kernels/list_kernels.cc
|
||
|
|
@@ -302,6 +302,10 @@ class TensorListReserve : public OpKernel {
|
||
|
|
PartialTensorShape element_shape;
|
||
|
|
OP_REQUIRES_OK(c, TensorShapeFromTensor(c->input(0), &element_shape));
|
||
|
|
int32 num_elements = c->input(1).scalar<int32>()();
|
||
|
|
+ OP_REQUIRES(c, num_elements >= 0,
|
||
|
|
+ errors::InvalidArgument("The num_elements to reserve must be a "
|
||
|
|
+ "non negative number, but got ",
|
||
|
|
+ num_elements));
|
||
|
|
TensorList output;
|
||
|
|
output.element_shape = element_shape;
|
||
|
|
output.element_dtype = element_dtype_;
|