Fix bugs in the indexing operator with latest NumPy (tensorflow#1895)

mrry · martinwicke · commit 1e2a17a445d8 · 2016-04-12T15:16:00.000-07:00
* Made the indexing operator accept objects that can be converted with `int()`.

Previously, the `_SliceHelper()` would test for `isinstance(..., int)`, which
is too restrictive.

* Replace `np.random.random_integers()` with `np.random.randint()`
diff --git a/tensorflow/python/kernel_tests/slice_op_test.py b/tensorflow/python/kernel_tests/slice_op_test.py
@@ -75,12 +75,15 @@ def _testSingleDimension(self, use_gpu):
       inp = np.random.rand(10).astype("f")
       a = tf.constant(inp, shape=[10], dtype=tf.float32)
 
-      hi = np.random.random_integers(0, 9)
+      hi = np.random.randint(0, 9)
       scalar_t = a[hi]
       scalar_val = scalar_t.eval()
       self.assertAllEqual(scalar_val, inp[hi])
 
-      lo = np.random.random_integers(0, hi)
+      if hi > 0:
+        lo = np.random.randint(0, hi)
+      else:
+        lo = 0
       slice_t = a[lo:hi]
       slice_val = slice_t.eval()
       self.assertAllEqual(slice_val, inp[lo:hi])
@@ -110,7 +113,7 @@ def _testIndexAndSlice(self, use_gpu):
       inp = np.random.rand(4, 4).astype("f")
       a = tf.constant(inp, shape=[4, 4], dtype=tf.float32)
 
-      x, y = np.random.random_integers(0, 3, size=2).tolist()
+      x, y = np.random.randint(0, 3, size=2).tolist()
       slice_t = a[x, 0:y]
       slice_val = slice_t.eval()
     self.assertAllEqual(slice_val, inp[x, 0:y])
@@ -142,9 +145,12 @@ def _testComplex(self, use_gpu):
       inp = np.random.rand(4, 10, 10, 4).astype("f")
       a = tf.constant(inp, dtype=tf.float32)
 
-      x = np.random.random_integers(0, 9)
-      z = np.random.random_integers(0, 9)
-      y = np.random.random_integers(0, z)
+      x = np.random.randint(0, 9)
+      z = np.random.randint(0, 9)
+      if z > 0:
+        y = np.random.randint(0, z)
+      else:
+        y = 0
       slice_t = a[:, x, y:z, :]
       self.assertAllEqual(slice_t.eval(), inp[:, x, y:z, :])
 
diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py
@@ -130,13 +130,7 @@ def _SliceHelper(tensor, slice_spec):
   sizes = []
   squeeze_dims = []
   for dim, s in enumerate(slice_spec):
-    if isinstance(s, int):
-      if s < 0:
-        raise NotImplementedError("Negative indices are currently unsupported")
-      indices.append(s)
-      sizes.append(1)
-      squeeze_dims.append(dim)
-    elif isinstance(s, _baseslice):
+    if isinstance(s, _baseslice):
       if s.step not in (None, 1):
         raise NotImplementedError(
             "Steps other than 1 are not currently supported")
@@ -161,7 +155,15 @@ def _SliceHelper(tensor, slice_spec):
     elif s is Ellipsis:
       raise NotImplementedError("Ellipsis is not currently supported")
     else:
-      raise TypeError("Bad slice index %s of type %s" % (s, type(s)))
+      try:
+        s = int(s)
+      except TypeError:
+        raise TypeError("Bad slice index %s of type %s" % (s, type(s)))
+      if s < 0:
+        raise NotImplementedError("Negative indices are currently unsupported")
+      indices.append(s)
+      sizes.append(1)
+      squeeze_dims.append(dim)
   sliced = slice(tensor, indices, sizes)
   if squeeze_dims:
     return squeeze(sliced, squeeze_dims=squeeze_dims)
