def testEinsum(self):
data1 = np.random.rand(3, 4, 5)
data2 = np.random.rand(4, 3, 2)
t1 = tensor(data1, chunk_size=2)
t2 = tensor(data2, chunk_size=3)
t = einsum('ijk, jil -> kl', t1, t2)
self.assertEqual(t.shape, (5, 2))
t = t.tiles()
self.assertEqual(len(t.chunks), 3)
# multiply(data1, data2)
data1 = np.random.rand(6, 6)
data2 = np.random.rand(6, 6)
t1 = tensor(data1, chunk_size=3)
t2 = tensor(data2, chunk_size=3)
t = einsum('..., ...', t1, t2)
self.assertEqual(t.shape, (6, 6))
t = t.tiles()
self.assertEqual(len(t.chunks), 4)
t = einsum('..., ...', t1, t2, optimize=True)
self.assertEqual(t.op.optimize, ['einsum_path', (0, 1)])
# test broadcast
data1 = np.random.rand(1, 10, 9)
data2 = np.random.rand(9, 6)
data3 = np.random.rand(10, 6)
data4 = np.random.rand(8,)
t1 = tensor(data1, chunk_size=(1, (5, 5), (3, 3, 3)))
t2 = tensor(data2, chunk_size=((3, 3, 3), (3, 3)))
t3 = tensor(data3, chunk_size=((6, 4), (4, 2)))
t4 = tensor(data4, chunk_size=3)
t = einsum('ajk,kl,jl,a->a', t1, t2, t3, t4, optimize='')
self.assertEqual(t.shape, (8,))
t = t.tiles()
self.assertEqual(len(t.chunks), 3)
def test_einsum():
data1 = np.random.rand(3, 4, 5)
data2 = np.random.rand(4, 3, 2)
t1 = tensor(data1, chunk_size=2)
t2 = tensor(data2, chunk_size=3)
t = einsum('ijk, jil -> kl', t1, t2)
assert t.shape == (5, 2)
t = tile(t)
assert len(t.chunks) == 3
# multiply(data1, data2)
data1 = np.random.rand(6, 6)
data2 = np.random.rand(6, 6)
t1 = tensor(data1, chunk_size=3)
t2 = tensor(data2, chunk_size=3)
t = einsum('..., ...', t1, t2)
assert t.shape == (6, 6)
t = tile(t)
assert len(t.chunks) == 4
t = einsum('..., ...', t1, t2, optimize=True)
assert t.op.optimize == ['einsum_path', (0, 1)]
# test broadcast
data1 = np.random.rand(1, 10, 9)
data2 = np.random.rand(9, 6)
data3 = np.random.rand(10, 6)
data4 = np.random.rand(8,)
t1 = tensor(data1, chunk_size=(1, (5, 5), (3, 3, 3)))
t2 = tensor(data2, chunk_size=((3, 3, 3), (3, 3)))
t3 = tensor(data3, chunk_size=((6, 4), (4, 2)))
t4 = tensor(data4, chunk_size=3)
t = einsum('ajk,kl,jl,a->a', t1, t2, t3, t4, optimize='')
assert t.shape == (8,)
t = tile(t)
assert len(t.chunks) == 3
def test_einsum_execution(setup):
data1 = np.random.rand(3, 4, 5)
data2 = np.random.rand(4, 3, 2)
t1 = tensor(data1, chunk_size=2)
t2 = tensor(data2, chunk_size=3)
t = einsum('ijk, jil -> kl', t1, t2)
res = t.execute().fetch()
expected = np.einsum('ijk, jil -> kl', data1, data2)
np.testing.assert_almost_equal(res, expected)
# dot
t = einsum('ijk, jil', t1, t2, optimize=True)
res = t.execute().fetch()
expected = np.einsum('ijk, jil', data1, data2, optimize=True)
np.testing.assert_almost_equal(res, expected)
# multiply(data1, data2)
data1 = np.random.rand(6, 6)
data2 = np.random.rand(6, 6)
t1 = tensor(data1, chunk_size=3)
t2 = tensor(data2, chunk_size=3)
t = einsum('..., ...', t1, t2, order='C')
res = t.execute().fetch()
expected = np.einsum('..., ...', data1, data2, order='C')
np.testing.assert_almost_equal(res, expected)
# sum(data, axis=-1)
data = np.random.rand(10)
t1 = tensor(data, chunk_size=3)
t = einsum('i->', t1, order='F')
res = t.execute().fetch()
expected = np.einsum('i->', data, order='F')
np.testing.assert_almost_equal(res, expected)
# sum(data, axis=0)
t1 = tensor(data)
t = einsum('...i->...', t1)
res = t.execute().fetch()
expected = np.einsum('...i->...', data)
np.testing.assert_almost_equal(res, expected)
# test broadcast
data1 = np.random.rand(1, 10, 9)
data2 = np.random.rand(9, 6)
data3 = np.random.rand(10, 6)
data4 = np.random.rand(8, )
t1 = tensor(data1, chunk_size=(1, (5, 5), (3, 3, 3)))
t2 = tensor(data2, chunk_size=((3, 3, 3), (3, 3)))
t3 = tensor(data3, chunk_size=((6, 4), (4, 2)))
t4 = tensor(data4, chunk_size=4)
t = einsum('ajk,kl,jl,a->a', t1, t2, t3, t4, optimize='optimal')
res = t.execute().fetch()
expected = np.einsum('ajk,kl,jl,a->a',
data1,
data2,
data3,
data4,
optimize='optimal')
np.testing.assert_almost_equal(res, expected)
t = einsum('ajk,kl,jl,a->a', t1, t2, t3, t4, optimize='greedy')
res = t.execute().fetch()
expected = np.einsum('ajk,kl,jl,a->a',
data1,
data2,
data3,
data4,
optimize='greedy')
np.testing.assert_almost_equal(res, expected)