Change StridedSlice to error on scalar input, in both

tensorflower-gardener · gunan · commit 083c4dc9edbf · 2016-10-22T22:25:18.000-07:00
the shape inference function and the kernel.
Change: 134434589
diff --git a/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc b/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc
@@ -72,7 +72,10 @@ REGISTER_OP("CudnnRNNParamsSize")
     .Attr(kRNNInputModeAttrs)
     .Attr(kRNNDirectionAttrs)
     .Output("params_size: S")
-    .SetShapeFn(shape_inference::ScalarShape)
+    .SetShapeFn([](InferenceContext* c) {
+      c->set_output(0, c->Vector(1));
+      return Status::OK();
+    })
     .Doc(strings::StrCat(R"doc(
 Return the params size that can be used by the Cudnn RNN model. Subsequent
 weight allocation and initialization should use this size.
diff --git a/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops_test.cc b/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops_test.cc
@@ -26,7 +26,7 @@ namespace tensorflow {
 
 TEST(CudnnRNNOpsTest, ParamsSize_ShapeFn) {
   ShapeInferenceTestOp op("CudnnRNNParamsSize");
-  INFER_OK(op, "[1];[1];[1]", "[]");
+  INFER_OK(op, "[1];[1];[1]", "[1]");
 }
 
 TEST(CudnnRNNOpsTest, ForwardLstm_ShapeFn) {
diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc
@@ -94,8 +94,8 @@ struct StridedSliceDenseSpec {
 }  // namespace
 
 template <class T>
-static void BuildDenseSpec(const StridedSliceSparseSpec& sparse,
-                           StridedSliceDenseSpec* dense) {
+static Status TF_MUST_USE_RESULT BuildDenseSpec(
+    const StridedSliceSparseSpec& sparse, StridedSliceDenseSpec* dense) {
   // Build expanded begin, end, strides, begin_mask, end_mask
   // to remove any ellipsis
   dense->begin.resize(dense->dims);
@@ -130,6 +130,12 @@ static void BuildDenseSpec(const StridedSliceSparseSpec& sparse,
       } else if ((1 << i) & sparse.new_axis_mask) {
         dense->final_shape_gather_indices.push_back(kNewAxis);
       } else {
+        if (full_index == dense->begin.size()) {
+          return errors::InvalidArgument("Index out of range using input dim ",
+                                         full_index, "; input has only ",
+                                         dense->dims, " dims");
+        }
+
         // Gather slicing spec into appropriate index
         dense->begin[full_index] = internal::SubtleMustCopy<T>(begin_flat(i));
         dense->end[full_index] = internal::SubtleMustCopy<T>(end_flat(i));
@@ -154,6 +160,7 @@ static void BuildDenseSpec(const StridedSliceSparseSpec& sparse,
       }
     }
   }
+  return Status::OK();
 }
 
 Status ValidateStridedSliceOp(
@@ -233,9 +240,9 @@ Status ValidateStridedSliceOp(
       input_shape.dims(), 0, 0, *begin, *end, *strides};
 
   if (begin_tensor.dtype() == DT_INT32) {
-    BuildDenseSpec<int32>(sparse_spec, &dense_spec);
+    TF_RETURN_IF_ERROR(BuildDenseSpec<int32>(sparse_spec, &dense_spec));
   } else if (begin_tensor.dtype() == DT_INT64) {
-    BuildDenseSpec<int64>(sparse_spec, &dense_spec);
+    TF_RETURN_IF_ERROR(BuildDenseSpec<int64>(sparse_spec, &dense_spec));
   } else {
     LOG(FATAL) << "begin must be either int32 or int64";
   }
diff --git a/tensorflow/python/kernel_tests/slice_op_test.py b/tensorflow/python/kernel_tests/slice_op_test.py
@@ -77,6 +77,34 @@ def testSingleDimension(self):
         slice_val = slice_t.eval()
         self.assertAllEqual(slice_val, inp[lo:hi])
 
+  def testScalarInput(self):
+    input_val = 0
+    with self.test_session() as sess:
+      # Test with constant input; shape inference fails.
+      with self.assertRaisesWithPredicateMatch(ValueError, "out of range"):
+        tf.constant(input_val)[:].get_shape()
+
+      # Test evaluating with non-constant input; kernel execution fails.
+      input_t = tf.placeholder(tf.int32)
+      slice_t = input_t[:]
+      with self.assertRaisesWithPredicateMatch(tf.errors.InvalidArgumentError,
+                                               "out of range"):
+        sess.run([slice_t], feed_dict={input_t: input_val})
+
+  def testInvalidIndex(self):
+    input_val = [1, 2]
+    with self.test_session() as sess:
+      # Test with constant input; shape inference fails.
+      with self.assertRaisesWithPredicateMatch(ValueError, "out of range"):
+        tf.constant(input_val)[1:, 1:].get_shape()
+
+      # Test evaluating with non-constant input; kernel execution fails.
+      input_t = tf.placeholder(tf.int32)
+      slice_t = input_t[1:, 1:]
+      with self.assertRaisesWithPredicateMatch(tf.errors.InvalidArgumentError,
+                                               "out of range"):
+        sess.run([slice_t], feed_dict={input_t: input_val})
+
   def _testSliceMatrixDim0(self, x, begin, size):
     with self.test_session(use_gpu=True):
       tf_ans = tf.slice(x, [begin, 0], [size, x.shape[1]]).eval()

Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,7 @@ namespace tensorflow {`
`26`	`26`
`27`	`27`	`TEST(CudnnRNNOpsTest, ParamsSize_ShapeFn) {`
`28`	`28`	`ShapeInferenceTestOp op("CudnnRNNParamsSize");`
`29`		`- INFER_OK(op, "[1];[1];[1]", "[]");`
	`29`	`+ INFER_OK(op, "[1];[1];[1]", "[1]");`
`30`	`30`	`}`
`31`	`31`
`32`	`32`	`TEST(CudnnRNNOpsTest, ForwardLstm_ShapeFn) {`