diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/array_ops_test.py')
| -rw-r--r-- | tensorflow/python/kernel_tests/array_ops_test.py | 52 |
1 files changed, 35 insertions, 17 deletions
diff --git a/tensorflow/python/kernel_tests/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops_test.py index 6eb9c66d06..76b80e60ea 100644 --- a/tensorflow/python/kernel_tests/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops_test.py @@ -107,22 +107,41 @@ class BooleanMaskTest(test_util.TensorFlowTestCase): def setUp(self): self.rng = np.random.RandomState(42) - def CheckVersusNumpy(self, ndims_mask, arr_shape, make_mask=None): + def CheckVersusNumpy(self, ndims_mask, arr_shape, make_mask=None, axis=None): """Check equivalence between boolean_mask and numpy masking.""" if make_mask is None: make_mask = lambda shape: self.rng.randint(0, 2, size=shape).astype(bool) arr = np.random.rand(*arr_shape) mask = make_mask(arr_shape[:ndims_mask]) - masked_arr = arr[mask] - with self.test_session(): - masked_tensor = array_ops.boolean_mask(arr, mask) + if axis is not None: + mask = make_mask(arr_shape[axis:ndims_mask+axis]) + if axis is None or axis == 0: + masked_arr = arr[mask] + elif axis == 1: + masked_arr = arr[:,mask] + elif axis == 2: + masked_arr = arr[:,:,mask] + with self.test_session() as sess: + masked_tensor = array_ops.boolean_mask(arr, mask, axis=axis) # Leading dimension size of masked_tensor is always unknown until runtime # since we don't how many elements will be kept. - self.assertAllEqual(masked_tensor.get_shape()[1:], masked_arr.shape[1:]) + leading = 1 if axis is None else axis + 1 + self.assertAllEqual(masked_tensor.get_shape()[leading:], + masked_arr.shape[leading:]) self.assertAllClose(masked_arr, masked_tensor.eval()) + def testMaskDim1ArrDim2Axis1(self): + ndims_mask = 1 + for arr_shape in [(1, 1), (2, 2), (2, 5)]: + self.CheckVersusNumpy(ndims_mask, arr_shape, axis=1) + + def testMaskDim2ArrDim2Axis1(self): + ndims_mask = 2 + for arr_shape in [(1, 1), (2, 2), (2, 5)]: + self.CheckVersusNumpy(ndims_mask, arr_shape, axis=1) + def testMaskDim1ArrDim1(self): ndims_mask = 1 for arr_shape in [(1,), (2,), (3,), (10,)]: @@ -486,7 +505,7 @@ class StridedSliceTest(test_util.TensorFlowTestCase): _ = checker2[...] _ = checker2[tuple()] - def testFloatSlicedArrayAndInt64IndicesGPU(self): + def testInt64GPU(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") with self.test_session(use_gpu=True, force_gpu=True): @@ -497,17 +516,6 @@ class StridedSliceTest(test_util.TensorFlowTestCase): s = array_ops.strided_slice(x, begin, end, strides) self.assertAllEqual([3.], self.evaluate(s)) - def testInt64SlicedArrayAndIndicesGPU(self): - if not test_util.is_gpu_available(): - self.skipTest("No GPU available") - with self.test_session(use_gpu=True, force_gpu=True): - x = constant_op.constant([1, 2, 3], dtype=dtypes.int64) - begin = constant_op.constant([2], dtype=dtypes.int64) - end = constant_op.constant([3], dtype=dtypes.int64) - strides = constant_op.constant([1], dtype=dtypes.int64) - s = array_ops.strided_slice(x, begin, end, strides) - self.assertAllEqual([3], self.evaluate(s)) - def testDegenerateSlices(self): with self.test_session(use_gpu=True): checker = StridedSliceChecker(self, StridedSliceChecker.REF_TENSOR) @@ -1070,6 +1078,16 @@ class PadTest(test_util.TensorFlowTestCase): [0, 0, 4, 5, 6, 0, 0], [0, 0, 0, 0, 0, 0, 0]]) +class InvertPermutationTest(test_util.TensorFlowTestCase): + + def testInvertPermutation(self): + for dtype in [dtypes.int32, dtypes.int64]: + with self.test_session(use_gpu=True): + x = constant_op.constant([3, 4, 0, 2, 1], dtype=dtype) + y = array_ops.invert_permutation(x) + self.assertAllEqual(y.get_shape(), [5]) + self.assertAllEqual(y.eval(), [2, 4, 3, 0, 1]) + if __name__ == "__main__": test_lib.main() |