def testIsIn(self): element = 2 * arange(4, chunk_size=1).reshape(2, 2) test_elements = [1, 2, 4, 8] mask = isin(element, test_elements) self.assertEqual(mask.shape, (2, 2)) self.assertEqual(mask.dtype, np.bool_) mask.tiles() self.assertEqual(len(mask.chunks), len(element.chunks)) self.assertEqual(len(mask.op.test_elements.chunks), 1) self.assertIs(mask.chunks[0].inputs[0], element.chunks[0].data) element = 2 * arange(4, chunk_size=1).reshape(2, 2) test_elements = tensor([1, 2, 4, 8], chunk_size=2) mask = isin(element, test_elements, invert=True) self.assertEqual(mask.shape, (2, 2)) self.assertEqual(mask.dtype, np.bool_) mask.tiles() self.assertEqual(len(mask.chunks), len(element.chunks)) self.assertEqual(len(mask.op.test_elements.chunks), 1) self.assertIs(mask.chunks[0].inputs[0], element.chunks[0].data) self.assertTrue(mask.chunks[0].op.invert)
def testIsInExecution(self): element = 2 * arange(4, chunk_size=1).reshape((2, 2)) test_elements = [1, 2, 4, 8] mask = isin(element, test_elements) res = self.executor.execute_tensor(mask, concat=True)[0] expected = np.isin(2 * np.arange(4).reshape((2, 2)), test_elements) np.testing.assert_equal(res, expected) res = self.executor.execute_tensor(element[mask], concat=True)[0] expected = np.array([2, 4]) np.testing.assert_equal(res, expected) mask = isin(element, test_elements, invert=True) res = self.executor.execute_tensor(mask, concat=True)[0] expected = np.isin(2 * np.arange(4).reshape((2, 2)), test_elements, invert=True) np.testing.assert_equal(res, expected) res = self.executor.execute_tensor(element[mask], concat=True)[0] expected = np.array([0, 6]) np.testing.assert_equal(res, expected) test_set = {1, 2, 4, 8} mask = isin(element, test_set) res = self.executor.execute_tensor(mask, concat=True)[0] expected = np.isin(2 * np.arange(4).reshape((2, 2)), test_set) np.testing.assert_equal(res, expected)