Fixes possible out-of-bounds access by strided slice.

PiperOrigin-RevId: 215269882
This commit is contained in:
Alexandre Passos 2018-10-01 13:28:17 -07:00 committed by TensorFlower Gardener
parent 3c6e6885f3
commit 1630584951
2 changed files with 9 additions and 1 deletions

View File

@ -149,7 +149,7 @@ class StridedSliceOp : public OpKernel {
// NDIM and T // NDIM and T
if (is_simple_slice && std::is_same<Device, CPUDevice>::value && if (is_simple_slice && std::is_same<Device, CPUDevice>::value &&
input_dims == 2 && processing_shape.dims() == 2 && input_dims == 2 && processing_shape.dims() == 2 &&
final_shape.dims() == 2) { final_shape.dims() == 2 && new_axis_mask == 0) {
MemCpyFunctor<T> functor; MemCpyFunctor<T> functor;
if (functor.Copy(input, begin, end, result)) { if (functor.Copy(input, begin, end, result)) {
return; return;

View File

@ -615,6 +615,14 @@ class StridedSliceTest(test_util.TensorFlowTestCase):
_ = checker[:, 0] _ = checker[:, 0]
_ = checker[:, :, 0] _ = checker[:, :, 0]
def testBothNewAxisAndShrink(self):
with self.test_session(use_gpu=True):
ones = array_ops.placeholder(shape=[2, 2], dtype=dtypes.int16)
self.assertAllEqual(
ones[array_ops.newaxis, :, 0].eval(
feed_dict={ones: [[1, 1], [1, 1]]}),
[[1, 1]])
def testTensorIndexing(self): def testTensorIndexing(self):
with self.test_session(use_gpu=True): with self.test_session(use_gpu=True):
raw = [[[[[1, 2, 4, 5], [5, 6, 7, 8], [9, 10, 11, 12]]], raw = [[[[[1, 2, 4, 5], [5, 6, 7, 8], [9, 10, 11, 12]]],