def testTakeExecution(self): data = np.random.rand(10, 20, 30) t = tensor(data, chunk_size=10) a = t.take([4, 1, 2, 6, 200]) res = self.executor.execute_tensor(a, concat=True)[0] expected = np.take(data, [4, 1, 2, 6, 200]) np.testing.assert_array_equal(res, expected) a = take(t, [5, 19, 2, 13], axis=1) res = self.executor.execute_tensor(a, concat=True)[0] expected = np.take(data, [5, 19, 2, 13], axis=1) np.testing.assert_array_equal(res, expected) with self.assertRaises(ValueError): take(t, [1, 3, 4], out=tensor(np.random.rand(4))) out = tensor([1, 2, 3, 4]) a = take(t, [4, 19, 2, 8], out=out) res = self.executor.execute_tensor(out, concat=True)[0] expected = np.take(data, [4, 19, 2, 8]) np.testing.assert_array_equal(res, expected)
def test_take_execution(setup): data = np.random.rand(10, 20, 30) t = tensor(data, chunk_size=10) a = t.take([4, 1, 2, 6, 200]) res = a.execute().fetch() expected = np.take(data, [4, 1, 2, 6, 200]) np.testing.assert_array_equal(res, expected) a = take(t, [5, 19, 2, 13], axis=1) res = a.execute().fetch() expected = np.take(data, [5, 19, 2, 13], axis=1) np.testing.assert_array_equal(res, expected) with pytest.raises(ValueError): take(t, [1, 3, 4], out=tensor(np.random.rand(4))) out = tensor([1, 2, 3, 4]) a = take(t, [4, 19, 2, 8], out=out) res = out.execute().fetch() expected = np.take(data, [4, 19, 2, 8]) np.testing.assert_array_equal(res, expected)
def testTakeExecution(self): data = np.random.rand(10, 20, 30) t = tensor(data, chunk_size=10) a = t.take([4, 1, 2, 6, 200]) res = self.executor.execute_tensor(a, concat=True)[0] expected = np.take(data, [4, 1, 2, 6, 200]) np.testing.assert_array_equal(res, expected) a = take(t, [5, 19, 2, 13], axis=1) res = self.executor.execute_tensor(a, concat=True)[0] expected = np.take(data, [5, 19, 2, 13], axis=1) np.testing.assert_array_equal(res, expected)