Fix scan ops for empty tensors

Change: 131302762
AshrafAamer · Aug 25, 2016 · b057e0d · b057e0d
1 parent a54aa42
commit b057e0d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/scan_ops.cc b/tensorflow/core/kernels/scan_ops.cc
@@ -62,6 +62,9 @@ class ScanOp : public OpKernel {
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &output));
 
+    // Exit early if there's nothing to compute
+    if (output_shape.num_elements() == 0) return;
+
     const Device& d = ctx->eigen_device<Device>();
     Reducer reducer;
 

diff --git a/tensorflow/python/kernel_tests/scan_ops_test.py b/tensorflow/python/kernel_tests/scan_ops_test.py
@@ -71,6 +71,11 @@ def _compareAll(self, x, axis):
       for reverse in [True, False]:
         self._compare(x, axis, exclusive, reverse)
 
+  def testEmpty(self):
+    for dtype in self.valid_dtypes:
+      x = np.zeros([0]).astype(dtype)
+      self._compareAll(x, 0)
+
   def test1D(self):
     for dtype in self.valid_dtypes:
       x = np.arange(1, 6).reshape([5]).astype(dtype)
@@ -155,6 +160,10 @@ def _compareAll(self, x, axis):
       for reverse in [True, False]:
         self._compare(x, axis, exclusive, reverse)
 
+  def testEmpty(self):
+    for dtype in self.valid_dtypes:
+      x = np.zeros([0]).astype(dtype)
+      self._compareAll(x, 0)
 
   def test1D(self):
     for dtype in self.valid_dtypes: