def testChooseExecution(self): options.tensor.chunk_size = 2 choices = [[0, 1, 2, 3], [10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33]] a = choose([2, 3, 1, 0], choices) res = self.executor.execute_tensor(a, concat=True)[0] expected = np.choose([2, 3, 1, 0], choices) np.testing.assert_array_equal(res, expected) a = choose([2, 4, 1, 0], choices, mode='clip') # 4 goes to 3 (4-1) expected = np.choose([2, 4, 1, 0], choices, mode='clip') res = self.executor.execute_tensor(a, concat=True)[0] np.testing.assert_array_equal(res, expected) a = choose([2, 4, 1, 0], choices, mode='wrap') # 4 goes to (4 mod 4) expected = np.choose([2, 4, 1, 0], choices, mode='wrap') # 4 goes to (4 mod 4) res = self.executor.execute_tensor(a, concat=True)[0] np.testing.assert_array_equal(res, expected) a = [[1, 0, 1], [0, 1, 0], [1, 0, 1]] choices = [-10, 10] b = choose(a, choices) expected = np.choose(a, choices) res = self.executor.execute_tensor(b, concat=True)[0] np.testing.assert_array_equal(res, expected) a = np.array([0, 1]).reshape((2, 1, 1)) c1 = np.array([1, 2, 3]).reshape((1, 3, 1)) c2 = np.array([-1, -2, -3, -4, -5]).reshape((1, 1, 5)) b = choose(a, (c1, c2)) expected = np.choose(a, (c1, c2)) res = self.executor.execute_tensor(b, concat=True)[0] np.testing.assert_array_equal(res, expected)
def testChoose(self): with option_context() as options: options.tensor.chunks = 2 choices = [[0, 1, 2, 3], [10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33]] a = choose([2, 3, 1, 0], choices) a.tiles() self.assertEqual(len(a.chunks), 2) self.assertIsInstance(a.chunks[0].op, type(a.op)) self.assertEqual(len(a.chunks[0].inputs), 5)