def testWhere(self): cond = tensor([[True, False], [False, True]], chunk_size=1) x = tensor([1, 2], chunk_size=1) y = tensor([3, 4], chunk_size=1) arr = where(cond, x, y) arr.tiles() self.assertEqual(len(arr.chunks), 4) self.assertTrue( np.array_equal(arr.chunks[0].inputs[0].op.data, [[True]])) self.assertTrue(np.array_equal(arr.chunks[0].inputs[1].op.data, [1])) self.assertTrue(np.array_equal(arr.chunks[0].inputs[2].op.data, [3])) self.assertTrue( np.array_equal(arr.chunks[1].inputs[0].op.data, [[False]])) self.assertTrue(np.array_equal(arr.chunks[1].inputs[1].op.data, [2])) self.assertTrue(np.array_equal(arr.chunks[1].inputs[2].op.data, [4])) self.assertTrue( np.array_equal(arr.chunks[2].inputs[0].op.data, [[False]])) self.assertTrue(np.array_equal(arr.chunks[2].inputs[1].op.data, [1])) self.assertTrue(np.array_equal(arr.chunks[2].inputs[2].op.data, [3])) self.assertTrue( np.array_equal(arr.chunks[3].inputs[0].op.data, [[True]])) self.assertTrue(np.array_equal(arr.chunks[3].inputs[1].op.data, [2])) self.assertTrue(np.array_equal(arr.chunks[3].inputs[2].op.data, [4])) with self.assertRaises(ValueError): where(cond, x) x = arange(9.).reshape(3, 3) y = where(x < 5, x, -1) self.assertEqual(y.dtype, np.float64)
def testWhereExecution(self): raw_cond = np.random.randint(0, 2, size=(4, 4), dtype='?') raw_x = np.random.rand(4, 1) raw_y = np.random.rand(4, 4) cond, x, y = tensor(raw_cond, chunk_size=2), tensor( raw_x, chunk_size=2), tensor(raw_y, chunk_size=2) arr = where(cond, x, y) res = self.executor.execute_tensor(arr, concat=True) self.assertTrue( np.array_equal(res[0], np.where(raw_cond, raw_x, raw_y))) raw_cond = sps.csr_matrix( np.random.randint(0, 2, size=(4, 4), dtype='?')) raw_x = sps.random(4, 1, density=.1) raw_y = sps.random(4, 4, density=.1) cond, x, y = tensor(raw_cond, chunk_size=2), tensor( raw_x, chunk_size=2), tensor(raw_y, chunk_size=2) arr = where(cond, x, y) res = self.executor.execute_tensor(arr, concat=True)[0] self.assertTrue( np.array_equal( res.toarray(), np.where(raw_cond.toarray(), raw_x.toarray(), raw_y.toarray())))