def test_einsum_misc(self): # This call used to crash because of a bug in # PyArray_AssignZero a = np.ones((1,2)) b = np.ones((2,2,1)) assert_equal(myeinsum('ij...,j...->i...',a,b), [[[2],[2]]]) # The iterator had an issue with buffering this reduction a = np.ones((5, 12, 4, 2, 3), np.int64) b = np.ones((5, 12, 11), np.int64) assert_equal(myeinsum('ijklm,ijn,ijn->',a,b,b), myeinsum('ijklm,ijn->',a,b)) # Issue #2027, was a problem in the contiguous 3-argument # inner loop implementation a = np.arange(1, 3) b = np.arange(1, 5).reshape(2, 2) c = np.arange(1, 9).reshape(4, 2) assert_equal(myeinsum('x,yx,zx->xzy', a, b, c), [[[1, 3], [3, 9], [5, 15], [7, 21]], [[8, 16], [16, 32], [24, 48], [32, 64]]])
def test_einsum_views(self): # pass-through a = np.arange(6) a.shape = (2,3) b = myeinsum("...", a) #assert_(b.base is a) assert_equal(b,a) ##b = myeinsum(a, [Ellipsis]) ##assert_(b.base is a) b = myeinsum("ij", a) #assert_(b.base is a) assert_equal(b, a) ##b = myeinsum(a, [0,1]) ##assert_(b.base is a) ##assert_equal(b, a) # transpose a = np.arange(6) a.shape = (2,3) b = myeinsum("ji", a) #assert_(b.base is a) assert_equal(b, a.T) ##b = myeinsum(a, [1,0]) ##assert_(b.base is a) ##assert_equal(b, a.T) # diagonal a = np.arange(9) a.shape = (3,3) b = myeinsum("ii->i", a) ##assert_(b.base is a) assert_equal(b, [a[i,i] for i in range(3)]) ##b = myeinsum(a, [0,0], [0]) ##assert_(b.base is a) ##assert_equal(b, [a[i,i] for i in range(3)]) # diagonal with various ways of broadcasting an additional dimension a = np.arange(27) a.shape = (3,3,3) b = myeinsum("...ii->...i", a) #assert_(b.base is a) assert_equal(b, [[x[i,i] for i in range(3)] for x in a]) ##b = myeinsum(a, [Ellipsis,0,0], [Ellipsis,0]) ##assert_(b.base is a) ##assert_equal(b, [[x[i,i] for i in range(3)] for x in a]) b = myeinsum("ii...->...i", a) #assert_(b.base is a) assert_equal(b, [[x[i,i] for i in range(3)] for x in a.transpose(2,0,1)]) ##b = myeinsum(a, [0,0,Ellipsis], [Ellipsis,0]) ##assert_(b.base is a) ##assert_equal(b, [[x[i,i] for i in range(3)] ## for x in a.transpose(2,0,1)]) b = myeinsum("...ii->i...", a) #assert_(b.base is a) assert_equal(b, [a[:,i,i] for i in range(3)]) ##b = myeinsum(a, [Ellipsis,0,0], [0,Ellipsis]) ##assert_(b.base is a) ##assert_equal(b, [a[:,i,i] for i in range(3)]) b = myeinsum("jii->ij", a) #assert_(b.base is a) assert_equal(b, [a[:,i,i] for i in range(3)]) ##b = myeinsum(a, [1,0,0], [0,1]) ##assert_(b.base is a) ##assert_equal(b, [a[:,i,i] for i in range(3)]) b = myeinsum("ii...->i...", a) #assert_(b.base is a) assert_equal(b, [a.transpose(2,0,1)[:,i,i] for i in range(3)]) ##b = myeinsum(a, [0,0,Ellipsis], [0,Ellipsis]) ##assert_(b.base is a) ##assert_equal(b, [a.transpose(2,0,1)[:,i,i] for i in range(3)]) b = myeinsum("i...i->i...", a) #assert_(b.base is a) assert_equal(b, [a.transpose(1,0,2)[:,i,i] for i in range(3)]) ##b = myeinsum(a, [0,Ellipsis,0], [0,Ellipsis]) ##assert_(b.base is a) ##assert_equal(b, [a.transpose(1,0,2)[:,i,i] for i in range(3)]) b = myeinsum("i...i->...i", a) #assert_(b.base is a) assert_equal(b, [[x[i,i] for i in range(3)] for x in a.transpose(1,0,2)]) ##b = myeinsum(a, [0,Ellipsis,0], [Ellipsis,0]) ##assert_(b.base is a) ##assert_equal(b, [[x[i,i] for i in range(3)] ## for x in a.transpose(1,0,2)]) # triple diagonal a = np.arange(27) a.shape = (3,3,3) b = myeinsum("iii->i", a) #assert_(b.base is a) assert_equal(b, [a[i,i,i] for i in range(3)]) ##b = myeinsum(a, [0,0,0], [0]) ##assert_(b.base is a) ##assert_equal(b, [a[i,i,i] for i in range(3)]) # swap axes a = np.arange(24) a.shape = (2,3,4) b = myeinsum("ijk->jik", a) #assert_(b.base is a) assert_equal(b, a.swapaxes(0,1))
def check_einsum_sums(self, dtype): # Check various sums. Does many sizes to exercise unrolled loops. # sum(a, axis=-1) for n in range(1,17): a = np.arange(n, dtype=dtype) assert_equal(myeinsum("i->", a), np.sum(a, axis=-1).astype(dtype)) ##assert_equal(myeinsum(a, [0], []), ## np.sum(a, axis=-1).astype(dtype)) for n in range(1,17): a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n) assert_equal(myeinsum("...i->...", a), np.sum(a, axis=-1).astype(dtype)) ##assert_equal(myeinsum(a, [Ellipsis,0], [Ellipsis]), ## np.sum(a, axis=-1).astype(dtype)) # sum(a, axis=0) for n in range(1,17): a = np.arange(2*n, dtype=dtype).reshape(2,n) assert_equal(myeinsum("i...->...", a), np.sum(a, axis=0).astype(dtype)) ##assert_equal(myeinsum(a, [0,Ellipsis], [Ellipsis]), ## np.sum(a, axis=0).astype(dtype)) for n in range(1,17): a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n) assert_equal(myeinsum("i...->...", a), np.sum(a, axis=0).astype(dtype)) ##assert_equal(myeinsum(a, [0,Ellipsis], [Ellipsis]), ## np.sum(a, axis=0).astype(dtype)) # trace(a) for n in range(1,17): a = np.arange(n*n, dtype=dtype).reshape(n,n) assert_equal(myeinsum("ii", a), np.trace(a).astype(dtype)) ##assert_equal(myeinsum(a, [0,0]), np.trace(a).astype(dtype)) # multiply(a, b) for n in range(1,17): a = np.arange(3*n, dtype=dtype).reshape(3,n) b = np.arange(2*3*n, dtype=dtype).reshape(2,3,n) assert_equal(myeinsum("..., ...", a, b), np.multiply(a, b)) ##assert_equal(myeinsum(a, [Ellipsis], b, [Ellipsis]), ## np.multiply(a, b)) # inner(a,b) for n in range(1,17): a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n) b = np.arange(n, dtype=dtype) assert_equal(myeinsum("...i, ...i", a, b), np.inner(a, b)) ##assert_equal(myeinsum(a, [Ellipsis,0], b, [Ellipsis,0]), ## np.inner(a, b)) for n in range(1,11): a = np.arange(n*3*2, dtype=dtype).reshape(n,3,2) b = np.arange(n, dtype=dtype) assert_equal(myeinsum("i..., i...", a, b), np.inner(a.T, b.T).T) ##assert_equal(myeinsum(a, [0,Ellipsis], b, [0,Ellipsis]), ## np.inner(a.T, b.T).T) # outer(a,b) for n in range(1,17): a = np.arange(3, dtype=dtype)+1 b = np.arange(n, dtype=dtype)+1 assert_equal(myeinsum("i,j", a, b), np.outer(a, b)) ##assert_equal(myeinsum(a, [0], b, [1]), np.outer(a, b)) # Suppress the complex warnings for the 'as f8' tests ctx = WarningManager() ctx.__enter__() try: warnings.simplefilter('ignore', np.ComplexWarning) # matvec(a,b) / a.dot(b) where a is matrix, b is vector for n in range(1,17): a = np.arange(4*n, dtype=dtype).reshape(4,n) b = np.arange(n, dtype=dtype) assert_equal(myeinsum("ij, j", a, b), np.dot(a, b)) ##assert_equal(myeinsum(a, [0,1], b, [1]), np.dot(a, b)) c = np.arange(4, dtype=dtype) myeinsum("ij,j", a, b, out=c, dtype='f8', casting='unsafe') assert_equal(c, np.dot(a.astype('f8'), b.astype('f8')).astype(dtype)) ##c[...] = 0 ##myeinsum(a, [0,1], b, [1], out=c, ## dtype='f8', casting='unsafe') ##assert_equal(c, ## np.dot(a.astype('f8'), ## b.astype('f8')).astype(dtype)) for n in range(1,17): a = np.arange(4*n, dtype=dtype).reshape(4,n) b = np.arange(n, dtype=dtype) assert_equal(myeinsum("ji,j", a.T, b.T), np.dot(b.T, a.T)) ##assert_equal(myeinsum(a.T, [1,0], b.T, [1]), np.dot(b.T, a.T)) c = np.arange(4, dtype=dtype) myeinsum("ji,j", a.T, b.T, out=c, dtype='f8', casting='unsafe') assert_equal(c, np.dot(b.T.astype('f8'), a.T.astype('f8')).astype(dtype)) ##c[...] = 0 ##myeinsum(a.T, [1,0], b.T, [1], out=c, ## dtype='f8', casting='unsafe') ##assert_equal(c, ## np.dot(b.T.astype('f8'), ## a.T.astype('f8')).astype(dtype)) # matmat(a,b) / a.dot(b) where a is matrix, b is matrix for n in range(1,17): if n < 8 or dtype != 'f2': a = np.arange(4*n, dtype=dtype).reshape(4,n) b = np.arange(n*6, dtype=dtype).reshape(n,6) assert_equal(myeinsum("ij,jk", a, b), np.dot(a, b)) ##assert_equal(myeinsum(a, [0,1], b, [1,2]), np.dot(a, b)) for n in range(1,17): a = np.arange(4*n, dtype=dtype).reshape(4,n) b = np.arange(n*6, dtype=dtype).reshape(n,6) c = np.arange(24, dtype=dtype).reshape(4,6) myeinsum("ij,jk", a, b, out=c, dtype='f8', casting='unsafe') assert_equal(c, np.dot(a.astype('f8'), b.astype('f8')).astype(dtype)) ##c[...] = 0 ##myeinsum(a, [0,1], b, [1,2], out=c, ## dtype='f8', casting='unsafe') ##assert_equal(c, ## np.dot(a.astype('f8'), ## b.astype('f8')).astype(dtype)) # matrix triple product (note this is not currently an efficient # way to multiply 3 matrices) a = np.arange(12, dtype=dtype).reshape(3,4) b = np.arange(20, dtype=dtype).reshape(4,5) c = np.arange(30, dtype=dtype).reshape(5,6) if dtype != 'f2': assert_equal(myeinsum("ij,jk,kl", a, b, c), a.dot(b).dot(c)) ##assert_equal(myeinsum(a, [0,1], b, [1,2], c, [2,3]), ## a.dot(b).dot(c)) d = np.arange(18, dtype=dtype).reshape(3,6) myeinsum("ij,jk,kl", a, b, c, out=d, dtype='f8', casting='unsafe') assert_equal(d, a.astype('f8').dot(b.astype('f8') ).dot(c.astype('f8')).astype(dtype)) ##d[...] = 0 ##myeinsum(a, [0,1], b, [1,2], c, [2,3], out=d, ## dtype='f8', casting='unsafe') ##assert_equal(d, a.astype('f8').dot(b.astype('f8') ## ).dot(c.astype('f8')).astype(dtype)) # tensordot(a, b) if np.dtype(dtype) != np.dtype('f2'): a = np.arange(60, dtype=dtype).reshape(3,4,5) b = np.arange(24, dtype=dtype).reshape(4,3,2) assert_equal(myeinsum("ijk, jil -> kl", a, b), np.tensordot(a,b, axes=([1,0],[0,1]))) ##assert_equal(myeinsum(a, [0,1,2], b, [1,0,3], [2,3]), ## np.tensordot(a,b, axes=([1,0],[0,1]))) c = np.arange(10, dtype=dtype).reshape(5,2) myeinsum("ijk,jil->kl", a, b, out=c, dtype='f8', casting='unsafe') assert_equal(c, np.tensordot(a.astype('f8'), b.astype('f8'), axes=([1,0],[0,1])).astype(dtype)) ##c[...] = 0 ##myeinsum(a, [0,1,2], b, [1,0,3], [2,3], out=c, ## dtype='f8', casting='unsafe') ##assert_equal(c, np.tensordot(a.astype('f8'), b.astype('f8'), ## axes=([1,0],[0,1])).astype(dtype)) finally: ctx.__exit__() """ # logical_and(logical_and(a!=0, b!=0), c!=0) a = np.array([1, 3, -2, 0, 12, 13, 0, 1], dtype=dtype) b = np.array([0, 3.5, 0., -2, 0, 1, 3, 12], dtype=dtype) c = np.array([True,True,False,True,True,False,True,True]) assert_equal(myeinsum("i,i,i->i", a, b, c, dtype='?', casting='unsafe'), np.logical_and(np.logical_and(a!=0, b!=0), c!=0)) ##assert_equal(myeinsum(a, [0], b, [0], c, [0], [0], ## dtype='?', casting='unsafe'), ## np.logical_and(np.logical_and(a!=0, b!=0), c!=0)) """ a = np.arange(9, dtype=dtype) assert_equal(myeinsum(",i->", 3, a), 3*np.sum(a)) ##assert_equal(myeinsum(3, [], a, [0], []), 3*np.sum(a)) assert_equal(myeinsum("i,->", a, 3), 3*np.sum(a)) ##assert_equal(myeinsum(a, [0], 3, [], []), 3*np.sum(a)) # Various stride0, contiguous, and SSE aligned variants for n in range(1,25): a = np.arange(n, dtype=dtype) if np.dtype(dtype).itemsize > 1: assert_equal(myeinsum("...,...",a,a), np.multiply(a,a)) assert_equal(myeinsum("i,i", a, a), np.dot(a,a)) assert_equal(myeinsum("i,->i", a, 2), 2*a) assert_equal(myeinsum(",i->i", 2, a), 2*a) assert_equal(myeinsum("i,->", a, 2), 2*np.sum(a)) assert_equal(myeinsum(",i->", 2, a), 2*np.sum(a)) assert_equal(myeinsum("...,...",a[1:],a[:-1]), np.multiply(a[1:],a[:-1])) assert_equal(myeinsum("i,i", a[1:], a[:-1]), np.dot(a[1:],a[:-1])) assert_equal(myeinsum("i,->i", a[1:], 2), 2*a[1:]) assert_equal(myeinsum(",i->i", 2, a[1:]), 2*a[1:]) assert_equal(myeinsum("i,->", a[1:], 2), 2*np.sum(a[1:])) assert_equal(myeinsum(",i->", 2, a[1:]), 2*np.sum(a[1:])) """ # An object array, summed as the data type a = np.arange(9, dtype=object) b = myeinsum("i->", a, dtype=dtype, casting='unsafe') assert_equal(b, np.sum(a)) assert_equal(b.dtype, np.dtype(dtype)) ##b = myeinsum(a, [0], [], dtype=dtype, casting='unsafe') ##assert_equal(b, np.sum(a)) ##assert_equal(b.dtype, np.dtype(dtype)) """ # A case which was failing (ticket #1885) p = np.arange(2) + 1 q = np.arange(4).reshape(2,2) + 3 r = np.arange(4).reshape(2,2) + 7 assert_equal(myeinsum('z,mz,zm->', p, q, r), 253)