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)
