test_einsum.py 49 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137
  1. import itertools
  2. import pytest
  3. import numpy as np
  4. from numpy.testing import (
  5. assert_, assert_equal, assert_array_equal, assert_almost_equal,
  6. assert_raises, suppress_warnings, assert_raises_regex, assert_allclose
  7. )
  8. # Setup for optimize einsum
  9. chars = 'abcdefghij'
  10. sizes = np.array([2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3])
  11. global_size_dict = dict(zip(chars, sizes))
  12. class TestEinsum:
  13. def test_einsum_errors(self):
  14. for do_opt in [True, False]:
  15. # Need enough arguments
  16. assert_raises(ValueError, np.einsum, optimize=do_opt)
  17. assert_raises(ValueError, np.einsum, "", optimize=do_opt)
  18. # subscripts must be a string
  19. assert_raises(TypeError, np.einsum, 0, 0, optimize=do_opt)
  20. # out parameter must be an array
  21. assert_raises(TypeError, np.einsum, "", 0, out='test',
  22. optimize=do_opt)
  23. # order parameter must be a valid order
  24. assert_raises(ValueError, np.einsum, "", 0, order='W',
  25. optimize=do_opt)
  26. # casting parameter must be a valid casting
  27. assert_raises(ValueError, np.einsum, "", 0, casting='blah',
  28. optimize=do_opt)
  29. # dtype parameter must be a valid dtype
  30. assert_raises(TypeError, np.einsum, "", 0, dtype='bad_data_type',
  31. optimize=do_opt)
  32. # other keyword arguments are rejected
  33. assert_raises(TypeError, np.einsum, "", 0, bad_arg=0,
  34. optimize=do_opt)
  35. # issue 4528 revealed a segfault with this call
  36. assert_raises(TypeError, np.einsum, *(None,)*63, optimize=do_opt)
  37. # number of operands must match count in subscripts string
  38. assert_raises(ValueError, np.einsum, "", 0, 0, optimize=do_opt)
  39. assert_raises(ValueError, np.einsum, ",", 0, [0], [0],
  40. optimize=do_opt)
  41. assert_raises(ValueError, np.einsum, ",", [0], optimize=do_opt)
  42. # can't have more subscripts than dimensions in the operand
  43. assert_raises(ValueError, np.einsum, "i", 0, optimize=do_opt)
  44. assert_raises(ValueError, np.einsum, "ij", [0, 0], optimize=do_opt)
  45. assert_raises(ValueError, np.einsum, "...i", 0, optimize=do_opt)
  46. assert_raises(ValueError, np.einsum, "i...j", [0, 0], optimize=do_opt)
  47. assert_raises(ValueError, np.einsum, "i...", 0, optimize=do_opt)
  48. assert_raises(ValueError, np.einsum, "ij...", [0, 0], optimize=do_opt)
  49. # invalid ellipsis
  50. assert_raises(ValueError, np.einsum, "i..", [0, 0], optimize=do_opt)
  51. assert_raises(ValueError, np.einsum, ".i...", [0, 0], optimize=do_opt)
  52. assert_raises(ValueError, np.einsum, "j->..j", [0, 0], optimize=do_opt)
  53. assert_raises(ValueError, np.einsum, "j->.j...", [0, 0], optimize=do_opt)
  54. # invalid subscript character
  55. assert_raises(ValueError, np.einsum, "i%...", [0, 0], optimize=do_opt)
  56. assert_raises(ValueError, np.einsum, "...j$", [0, 0], optimize=do_opt)
  57. assert_raises(ValueError, np.einsum, "i->&", [0, 0], optimize=do_opt)
  58. # output subscripts must appear in input
  59. assert_raises(ValueError, np.einsum, "i->ij", [0, 0], optimize=do_opt)
  60. # output subscripts may only be specified once
  61. assert_raises(ValueError, np.einsum, "ij->jij", [[0, 0], [0, 0]],
  62. optimize=do_opt)
  63. # dimensions much match when being collapsed
  64. assert_raises(ValueError, np.einsum, "ii",
  65. np.arange(6).reshape(2, 3), optimize=do_opt)
  66. assert_raises(ValueError, np.einsum, "ii->i",
  67. np.arange(6).reshape(2, 3), optimize=do_opt)
  68. # broadcasting to new dimensions must be enabled explicitly
  69. assert_raises(ValueError, np.einsum, "i", np.arange(6).reshape(2, 3),
  70. optimize=do_opt)
  71. assert_raises(ValueError, np.einsum, "i->i", [[0, 1], [0, 1]],
  72. out=np.arange(4).reshape(2, 2), optimize=do_opt)
  73. with assert_raises_regex(ValueError, "'b'"):
  74. # gh-11221 - 'c' erroneously appeared in the error message
  75. a = np.ones((3, 3, 4, 5, 6))
  76. b = np.ones((3, 4, 5))
  77. np.einsum('aabcb,abc', a, b)
  78. # Check order kwarg, asanyarray allows 1d to pass through
  79. assert_raises(ValueError, np.einsum, "i->i", np.arange(6).reshape(-1, 1),
  80. optimize=do_opt, order='d')
  81. def test_einsum_views(self):
  82. # pass-through
  83. for do_opt in [True, False]:
  84. a = np.arange(6)
  85. a.shape = (2, 3)
  86. b = np.einsum("...", a, optimize=do_opt)
  87. assert_(b.base is a)
  88. b = np.einsum(a, [Ellipsis], optimize=do_opt)
  89. assert_(b.base is a)
  90. b = np.einsum("ij", a, optimize=do_opt)
  91. assert_(b.base is a)
  92. assert_equal(b, a)
  93. b = np.einsum(a, [0, 1], optimize=do_opt)
  94. assert_(b.base is a)
  95. assert_equal(b, a)
  96. # output is writeable whenever input is writeable
  97. b = np.einsum("...", a, optimize=do_opt)
  98. assert_(b.flags['WRITEABLE'])
  99. a.flags['WRITEABLE'] = False
  100. b = np.einsum("...", a, optimize=do_opt)
  101. assert_(not b.flags['WRITEABLE'])
  102. # transpose
  103. a = np.arange(6)
  104. a.shape = (2, 3)
  105. b = np.einsum("ji", a, optimize=do_opt)
  106. assert_(b.base is a)
  107. assert_equal(b, a.T)
  108. b = np.einsum(a, [1, 0], optimize=do_opt)
  109. assert_(b.base is a)
  110. assert_equal(b, a.T)
  111. # diagonal
  112. a = np.arange(9)
  113. a.shape = (3, 3)
  114. b = np.einsum("ii->i", a, optimize=do_opt)
  115. assert_(b.base is a)
  116. assert_equal(b, [a[i, i] for i in range(3)])
  117. b = np.einsum(a, [0, 0], [0], optimize=do_opt)
  118. assert_(b.base is a)
  119. assert_equal(b, [a[i, i] for i in range(3)])
  120. # diagonal with various ways of broadcasting an additional dimension
  121. a = np.arange(27)
  122. a.shape = (3, 3, 3)
  123. b = np.einsum("...ii->...i", a, optimize=do_opt)
  124. assert_(b.base is a)
  125. assert_equal(b, [[x[i, i] for i in range(3)] for x in a])
  126. b = np.einsum(a, [Ellipsis, 0, 0], [Ellipsis, 0], optimize=do_opt)
  127. assert_(b.base is a)
  128. assert_equal(b, [[x[i, i] for i in range(3)] for x in a])
  129. b = np.einsum("ii...->...i", a, optimize=do_opt)
  130. assert_(b.base is a)
  131. assert_equal(b, [[x[i, i] for i in range(3)]
  132. for x in a.transpose(2, 0, 1)])
  133. b = np.einsum(a, [0, 0, Ellipsis], [Ellipsis, 0], optimize=do_opt)
  134. assert_(b.base is a)
  135. assert_equal(b, [[x[i, i] for i in range(3)]
  136. for x in a.transpose(2, 0, 1)])
  137. b = np.einsum("...ii->i...", a, optimize=do_opt)
  138. assert_(b.base is a)
  139. assert_equal(b, [a[:, i, i] for i in range(3)])
  140. b = np.einsum(a, [Ellipsis, 0, 0], [0, Ellipsis], optimize=do_opt)
  141. assert_(b.base is a)
  142. assert_equal(b, [a[:, i, i] for i in range(3)])
  143. b = np.einsum("jii->ij", a, optimize=do_opt)
  144. assert_(b.base is a)
  145. assert_equal(b, [a[:, i, i] for i in range(3)])
  146. b = np.einsum(a, [1, 0, 0], [0, 1], optimize=do_opt)
  147. assert_(b.base is a)
  148. assert_equal(b, [a[:, i, i] for i in range(3)])
  149. b = np.einsum("ii...->i...", a, optimize=do_opt)
  150. assert_(b.base is a)
  151. assert_equal(b, [a.transpose(2, 0, 1)[:, i, i] for i in range(3)])
  152. b = np.einsum(a, [0, 0, Ellipsis], [0, Ellipsis], optimize=do_opt)
  153. assert_(b.base is a)
  154. assert_equal(b, [a.transpose(2, 0, 1)[:, i, i] for i in range(3)])
  155. b = np.einsum("i...i->i...", a, optimize=do_opt)
  156. assert_(b.base is a)
  157. assert_equal(b, [a.transpose(1, 0, 2)[:, i, i] for i in range(3)])
  158. b = np.einsum(a, [0, Ellipsis, 0], [0, Ellipsis], optimize=do_opt)
  159. assert_(b.base is a)
  160. assert_equal(b, [a.transpose(1, 0, 2)[:, i, i] for i in range(3)])
  161. b = np.einsum("i...i->...i", a, optimize=do_opt)
  162. assert_(b.base is a)
  163. assert_equal(b, [[x[i, i] for i in range(3)]
  164. for x in a.transpose(1, 0, 2)])
  165. b = np.einsum(a, [0, Ellipsis, 0], [Ellipsis, 0], optimize=do_opt)
  166. assert_(b.base is a)
  167. assert_equal(b, [[x[i, i] for i in range(3)]
  168. for x in a.transpose(1, 0, 2)])
  169. # triple diagonal
  170. a = np.arange(27)
  171. a.shape = (3, 3, 3)
  172. b = np.einsum("iii->i", a, optimize=do_opt)
  173. assert_(b.base is a)
  174. assert_equal(b, [a[i, i, i] for i in range(3)])
  175. b = np.einsum(a, [0, 0, 0], [0], optimize=do_opt)
  176. assert_(b.base is a)
  177. assert_equal(b, [a[i, i, i] for i in range(3)])
  178. # swap axes
  179. a = np.arange(24)
  180. a.shape = (2, 3, 4)
  181. b = np.einsum("ijk->jik", a, optimize=do_opt)
  182. assert_(b.base is a)
  183. assert_equal(b, a.swapaxes(0, 1))
  184. b = np.einsum(a, [0, 1, 2], [1, 0, 2], optimize=do_opt)
  185. assert_(b.base is a)
  186. assert_equal(b, a.swapaxes(0, 1))
  187. @np._no_nep50_warning()
  188. def check_einsum_sums(self, dtype, do_opt=False):
  189. dtype = np.dtype(dtype)
  190. # Check various sums. Does many sizes to exercise unrolled loops.
  191. # sum(a, axis=-1)
  192. for n in range(1, 17):
  193. a = np.arange(n, dtype=dtype)
  194. assert_equal(np.einsum("i->", a, optimize=do_opt),
  195. np.sum(a, axis=-1).astype(dtype))
  196. assert_equal(np.einsum(a, [0], [], optimize=do_opt),
  197. np.sum(a, axis=-1).astype(dtype))
  198. for n in range(1, 17):
  199. a = np.arange(2*3*n, dtype=dtype).reshape(2, 3, n)
  200. assert_equal(np.einsum("...i->...", a, optimize=do_opt),
  201. np.sum(a, axis=-1).astype(dtype))
  202. assert_equal(np.einsum(a, [Ellipsis, 0], [Ellipsis], optimize=do_opt),
  203. np.sum(a, axis=-1).astype(dtype))
  204. # sum(a, axis=0)
  205. for n in range(1, 17):
  206. a = np.arange(2*n, dtype=dtype).reshape(2, n)
  207. assert_equal(np.einsum("i...->...", a, optimize=do_opt),
  208. np.sum(a, axis=0).astype(dtype))
  209. assert_equal(np.einsum(a, [0, Ellipsis], [Ellipsis], optimize=do_opt),
  210. np.sum(a, axis=0).astype(dtype))
  211. for n in range(1, 17):
  212. a = np.arange(2*3*n, dtype=dtype).reshape(2, 3, n)
  213. assert_equal(np.einsum("i...->...", a, optimize=do_opt),
  214. np.sum(a, axis=0).astype(dtype))
  215. assert_equal(np.einsum(a, [0, Ellipsis], [Ellipsis], optimize=do_opt),
  216. np.sum(a, axis=0).astype(dtype))
  217. # trace(a)
  218. for n in range(1, 17):
  219. a = np.arange(n*n, dtype=dtype).reshape(n, n)
  220. assert_equal(np.einsum("ii", a, optimize=do_opt),
  221. np.trace(a).astype(dtype))
  222. assert_equal(np.einsum(a, [0, 0], optimize=do_opt),
  223. np.trace(a).astype(dtype))
  224. # gh-15961: should accept numpy int64 type in subscript list
  225. np_array = np.asarray([0, 0])
  226. assert_equal(np.einsum(a, np_array, optimize=do_opt),
  227. np.trace(a).astype(dtype))
  228. assert_equal(np.einsum(a, list(np_array), optimize=do_opt),
  229. np.trace(a).astype(dtype))
  230. # multiply(a, b)
  231. assert_equal(np.einsum("..., ...", 3, 4), 12) # scalar case
  232. for n in range(1, 17):
  233. a = np.arange(3 * n, dtype=dtype).reshape(3, n)
  234. b = np.arange(2 * 3 * n, dtype=dtype).reshape(2, 3, n)
  235. assert_equal(np.einsum("..., ...", a, b, optimize=do_opt),
  236. np.multiply(a, b))
  237. assert_equal(np.einsum(a, [Ellipsis], b, [Ellipsis], optimize=do_opt),
  238. np.multiply(a, b))
  239. # inner(a,b)
  240. for n in range(1, 17):
  241. a = np.arange(2 * 3 * n, dtype=dtype).reshape(2, 3, n)
  242. b = np.arange(n, dtype=dtype)
  243. assert_equal(np.einsum("...i, ...i", a, b, optimize=do_opt), np.inner(a, b))
  244. assert_equal(np.einsum(a, [Ellipsis, 0], b, [Ellipsis, 0], optimize=do_opt),
  245. np.inner(a, b))
  246. for n in range(1, 11):
  247. a = np.arange(n * 3 * 2, dtype=dtype).reshape(n, 3, 2)
  248. b = np.arange(n, dtype=dtype)
  249. assert_equal(np.einsum("i..., i...", a, b, optimize=do_opt),
  250. np.inner(a.T, b.T).T)
  251. assert_equal(np.einsum(a, [0, Ellipsis], b, [0, Ellipsis], optimize=do_opt),
  252. np.inner(a.T, b.T).T)
  253. # outer(a,b)
  254. for n in range(1, 17):
  255. a = np.arange(3, dtype=dtype)+1
  256. b = np.arange(n, dtype=dtype)+1
  257. assert_equal(np.einsum("i,j", a, b, optimize=do_opt),
  258. np.outer(a, b))
  259. assert_equal(np.einsum(a, [0], b, [1], optimize=do_opt),
  260. np.outer(a, b))
  261. # Suppress the complex warnings for the 'as f8' tests
  262. with suppress_warnings() as sup:
  263. sup.filter(np.ComplexWarning)
  264. # matvec(a,b) / a.dot(b) where a is matrix, b is vector
  265. for n in range(1, 17):
  266. a = np.arange(4*n, dtype=dtype).reshape(4, n)
  267. b = np.arange(n, dtype=dtype)
  268. assert_equal(np.einsum("ij, j", a, b, optimize=do_opt),
  269. np.dot(a, b))
  270. assert_equal(np.einsum(a, [0, 1], b, [1], optimize=do_opt),
  271. np.dot(a, b))
  272. c = np.arange(4, dtype=dtype)
  273. np.einsum("ij,j", a, b, out=c,
  274. dtype='f8', casting='unsafe', optimize=do_opt)
  275. assert_equal(c,
  276. np.dot(a.astype('f8'),
  277. b.astype('f8')).astype(dtype))
  278. c[...] = 0
  279. np.einsum(a, [0, 1], b, [1], out=c,
  280. dtype='f8', casting='unsafe', optimize=do_opt)
  281. assert_equal(c,
  282. np.dot(a.astype('f8'),
  283. b.astype('f8')).astype(dtype))
  284. for n in range(1, 17):
  285. a = np.arange(4*n, dtype=dtype).reshape(4, n)
  286. b = np.arange(n, dtype=dtype)
  287. assert_equal(np.einsum("ji,j", a.T, b.T, optimize=do_opt),
  288. np.dot(b.T, a.T))
  289. assert_equal(np.einsum(a.T, [1, 0], b.T, [1], optimize=do_opt),
  290. np.dot(b.T, a.T))
  291. c = np.arange(4, dtype=dtype)
  292. np.einsum("ji,j", a.T, b.T, out=c,
  293. dtype='f8', casting='unsafe', optimize=do_opt)
  294. assert_equal(c,
  295. np.dot(b.T.astype('f8'),
  296. a.T.astype('f8')).astype(dtype))
  297. c[...] = 0
  298. np.einsum(a.T, [1, 0], b.T, [1], out=c,
  299. dtype='f8', casting='unsafe', optimize=do_opt)
  300. assert_equal(c,
  301. np.dot(b.T.astype('f8'),
  302. a.T.astype('f8')).astype(dtype))
  303. # matmat(a,b) / a.dot(b) where a is matrix, b is matrix
  304. for n in range(1, 17):
  305. if n < 8 or dtype != 'f2':
  306. a = np.arange(4*n, dtype=dtype).reshape(4, n)
  307. b = np.arange(n*6, dtype=dtype).reshape(n, 6)
  308. assert_equal(np.einsum("ij,jk", a, b, optimize=do_opt),
  309. np.dot(a, b))
  310. assert_equal(np.einsum(a, [0, 1], b, [1, 2], optimize=do_opt),
  311. np.dot(a, b))
  312. for n in range(1, 17):
  313. a = np.arange(4*n, dtype=dtype).reshape(4, n)
  314. b = np.arange(n*6, dtype=dtype).reshape(n, 6)
  315. c = np.arange(24, dtype=dtype).reshape(4, 6)
  316. np.einsum("ij,jk", a, b, out=c, dtype='f8', casting='unsafe',
  317. optimize=do_opt)
  318. assert_equal(c,
  319. np.dot(a.astype('f8'),
  320. b.astype('f8')).astype(dtype))
  321. c[...] = 0
  322. np.einsum(a, [0, 1], b, [1, 2], out=c,
  323. dtype='f8', casting='unsafe', optimize=do_opt)
  324. assert_equal(c,
  325. np.dot(a.astype('f8'),
  326. b.astype('f8')).astype(dtype))
  327. # matrix triple product (note this is not currently an efficient
  328. # way to multiply 3 matrices)
  329. a = np.arange(12, dtype=dtype).reshape(3, 4)
  330. b = np.arange(20, dtype=dtype).reshape(4, 5)
  331. c = np.arange(30, dtype=dtype).reshape(5, 6)
  332. if dtype != 'f2':
  333. assert_equal(np.einsum("ij,jk,kl", a, b, c, optimize=do_opt),
  334. a.dot(b).dot(c))
  335. assert_equal(np.einsum(a, [0, 1], b, [1, 2], c, [2, 3],
  336. optimize=do_opt), a.dot(b).dot(c))
  337. d = np.arange(18, dtype=dtype).reshape(3, 6)
  338. np.einsum("ij,jk,kl", a, b, c, out=d,
  339. dtype='f8', casting='unsafe', optimize=do_opt)
  340. tgt = a.astype('f8').dot(b.astype('f8'))
  341. tgt = tgt.dot(c.astype('f8')).astype(dtype)
  342. assert_equal(d, tgt)
  343. d[...] = 0
  344. np.einsum(a, [0, 1], b, [1, 2], c, [2, 3], out=d,
  345. dtype='f8', casting='unsafe', optimize=do_opt)
  346. tgt = a.astype('f8').dot(b.astype('f8'))
  347. tgt = tgt.dot(c.astype('f8')).astype(dtype)
  348. assert_equal(d, tgt)
  349. # tensordot(a, b)
  350. if np.dtype(dtype) != np.dtype('f2'):
  351. a = np.arange(60, dtype=dtype).reshape(3, 4, 5)
  352. b = np.arange(24, dtype=dtype).reshape(4, 3, 2)
  353. assert_equal(np.einsum("ijk, jil -> kl", a, b),
  354. np.tensordot(a, b, axes=([1, 0], [0, 1])))
  355. assert_equal(np.einsum(a, [0, 1, 2], b, [1, 0, 3], [2, 3]),
  356. np.tensordot(a, b, axes=([1, 0], [0, 1])))
  357. c = np.arange(10, dtype=dtype).reshape(5, 2)
  358. np.einsum("ijk,jil->kl", a, b, out=c,
  359. dtype='f8', casting='unsafe', optimize=do_opt)
  360. assert_equal(c, np.tensordot(a.astype('f8'), b.astype('f8'),
  361. axes=([1, 0], [0, 1])).astype(dtype))
  362. c[...] = 0
  363. np.einsum(a, [0, 1, 2], b, [1, 0, 3], [2, 3], out=c,
  364. dtype='f8', casting='unsafe', optimize=do_opt)
  365. assert_equal(c, np.tensordot(a.astype('f8'), b.astype('f8'),
  366. axes=([1, 0], [0, 1])).astype(dtype))
  367. # logical_and(logical_and(a!=0, b!=0), c!=0)
  368. neg_val = -2 if dtype.kind != "u" else np.iinfo(dtype).max - 1
  369. a = np.array([1, 3, neg_val, 0, 12, 13, 0, 1], dtype=dtype)
  370. b = np.array([0, 3.5, 0., neg_val, 0, 1, 3, 12], dtype=dtype)
  371. c = np.array([True, True, False, True, True, False, True, True])
  372. assert_equal(np.einsum("i,i,i->i", a, b, c,
  373. dtype='?', casting='unsafe', optimize=do_opt),
  374. np.logical_and(np.logical_and(a != 0, b != 0), c != 0))
  375. assert_equal(np.einsum(a, [0], b, [0], c, [0], [0],
  376. dtype='?', casting='unsafe'),
  377. np.logical_and(np.logical_and(a != 0, b != 0), c != 0))
  378. a = np.arange(9, dtype=dtype)
  379. assert_equal(np.einsum(",i->", 3, a), 3*np.sum(a))
  380. assert_equal(np.einsum(3, [], a, [0], []), 3*np.sum(a))
  381. assert_equal(np.einsum("i,->", a, 3), 3*np.sum(a))
  382. assert_equal(np.einsum(a, [0], 3, [], []), 3*np.sum(a))
  383. # Various stride0, contiguous, and SSE aligned variants
  384. for n in range(1, 25):
  385. a = np.arange(n, dtype=dtype)
  386. if np.dtype(dtype).itemsize > 1:
  387. assert_equal(np.einsum("...,...", a, a, optimize=do_opt),
  388. np.multiply(a, a))
  389. assert_equal(np.einsum("i,i", a, a, optimize=do_opt), np.dot(a, a))
  390. assert_equal(np.einsum("i,->i", a, 2, optimize=do_opt), 2*a)
  391. assert_equal(np.einsum(",i->i", 2, a, optimize=do_opt), 2*a)
  392. assert_equal(np.einsum("i,->", a, 2, optimize=do_opt), 2*np.sum(a))
  393. assert_equal(np.einsum(",i->", 2, a, optimize=do_opt), 2*np.sum(a))
  394. assert_equal(np.einsum("...,...", a[1:], a[:-1], optimize=do_opt),
  395. np.multiply(a[1:], a[:-1]))
  396. assert_equal(np.einsum("i,i", a[1:], a[:-1], optimize=do_opt),
  397. np.dot(a[1:], a[:-1]))
  398. assert_equal(np.einsum("i,->i", a[1:], 2, optimize=do_opt), 2*a[1:])
  399. assert_equal(np.einsum(",i->i", 2, a[1:], optimize=do_opt), 2*a[1:])
  400. assert_equal(np.einsum("i,->", a[1:], 2, optimize=do_opt),
  401. 2*np.sum(a[1:]))
  402. assert_equal(np.einsum(",i->", 2, a[1:], optimize=do_opt),
  403. 2*np.sum(a[1:]))
  404. # An object array, summed as the data type
  405. a = np.arange(9, dtype=object)
  406. b = np.einsum("i->", a, dtype=dtype, casting='unsafe')
  407. assert_equal(b, np.sum(a))
  408. assert_equal(b.dtype, np.dtype(dtype))
  409. b = np.einsum(a, [0], [], dtype=dtype, casting='unsafe')
  410. assert_equal(b, np.sum(a))
  411. assert_equal(b.dtype, np.dtype(dtype))
  412. # A case which was failing (ticket #1885)
  413. p = np.arange(2) + 1
  414. q = np.arange(4).reshape(2, 2) + 3
  415. r = np.arange(4).reshape(2, 2) + 7
  416. assert_equal(np.einsum('z,mz,zm->', p, q, r), 253)
  417. # singleton dimensions broadcast (gh-10343)
  418. p = np.ones((10,2))
  419. q = np.ones((1,2))
  420. assert_array_equal(np.einsum('ij,ij->j', p, q, optimize=True),
  421. np.einsum('ij,ij->j', p, q, optimize=False))
  422. assert_array_equal(np.einsum('ij,ij->j', p, q, optimize=True),
  423. [10.] * 2)
  424. # a blas-compatible contraction broadcasting case which was failing
  425. # for optimize=True (ticket #10930)
  426. x = np.array([2., 3.])
  427. y = np.array([4.])
  428. assert_array_equal(np.einsum("i, i", x, y, optimize=False), 20.)
  429. assert_array_equal(np.einsum("i, i", x, y, optimize=True), 20.)
  430. # all-ones array was bypassing bug (ticket #10930)
  431. p = np.ones((1, 5)) / 2
  432. q = np.ones((5, 5)) / 2
  433. for optimize in (True, False):
  434. assert_array_equal(np.einsum("...ij,...jk->...ik", p, p,
  435. optimize=optimize),
  436. np.einsum("...ij,...jk->...ik", p, q,
  437. optimize=optimize))
  438. assert_array_equal(np.einsum("...ij,...jk->...ik", p, q,
  439. optimize=optimize),
  440. np.full((1, 5), 1.25))
  441. # Cases which were failing (gh-10899)
  442. x = np.eye(2, dtype=dtype)
  443. y = np.ones(2, dtype=dtype)
  444. assert_array_equal(np.einsum("ji,i->", x, y, optimize=optimize),
  445. [2.]) # contig_contig_outstride0_two
  446. assert_array_equal(np.einsum("i,ij->", y, x, optimize=optimize),
  447. [2.]) # stride0_contig_outstride0_two
  448. assert_array_equal(np.einsum("ij,i->", x, y, optimize=optimize),
  449. [2.]) # contig_stride0_outstride0_two
  450. def test_einsum_sums_int8(self):
  451. self.check_einsum_sums('i1')
  452. def test_einsum_sums_uint8(self):
  453. self.check_einsum_sums('u1')
  454. def test_einsum_sums_int16(self):
  455. self.check_einsum_sums('i2')
  456. def test_einsum_sums_uint16(self):
  457. self.check_einsum_sums('u2')
  458. def test_einsum_sums_int32(self):
  459. self.check_einsum_sums('i4')
  460. self.check_einsum_sums('i4', True)
  461. def test_einsum_sums_uint32(self):
  462. self.check_einsum_sums('u4')
  463. self.check_einsum_sums('u4', True)
  464. def test_einsum_sums_int64(self):
  465. self.check_einsum_sums('i8')
  466. def test_einsum_sums_uint64(self):
  467. self.check_einsum_sums('u8')
  468. def test_einsum_sums_float16(self):
  469. self.check_einsum_sums('f2')
  470. def test_einsum_sums_float32(self):
  471. self.check_einsum_sums('f4')
  472. def test_einsum_sums_float64(self):
  473. self.check_einsum_sums('f8')
  474. self.check_einsum_sums('f8', True)
  475. def test_einsum_sums_longdouble(self):
  476. self.check_einsum_sums(np.longdouble)
  477. def test_einsum_sums_cfloat64(self):
  478. self.check_einsum_sums('c8')
  479. self.check_einsum_sums('c8', True)
  480. def test_einsum_sums_cfloat128(self):
  481. self.check_einsum_sums('c16')
  482. def test_einsum_sums_clongdouble(self):
  483. self.check_einsum_sums(np.clongdouble)
  484. def test_einsum_misc(self):
  485. # This call used to crash because of a bug in
  486. # PyArray_AssignZero
  487. a = np.ones((1, 2))
  488. b = np.ones((2, 2, 1))
  489. assert_equal(np.einsum('ij...,j...->i...', a, b), [[[2], [2]]])
  490. assert_equal(np.einsum('ij...,j...->i...', a, b, optimize=True), [[[2], [2]]])
  491. # Regression test for issue #10369 (test unicode inputs with Python 2)
  492. assert_equal(np.einsum('ij...,j...->i...', a, b), [[[2], [2]]])
  493. assert_equal(np.einsum('...i,...i', [1, 2, 3], [2, 3, 4]), 20)
  494. assert_equal(np.einsum('...i,...i', [1, 2, 3], [2, 3, 4],
  495. optimize='greedy'), 20)
  496. # The iterator had an issue with buffering this reduction
  497. a = np.ones((5, 12, 4, 2, 3), np.int64)
  498. b = np.ones((5, 12, 11), np.int64)
  499. assert_equal(np.einsum('ijklm,ijn,ijn->', a, b, b),
  500. np.einsum('ijklm,ijn->', a, b))
  501. assert_equal(np.einsum('ijklm,ijn,ijn->', a, b, b, optimize=True),
  502. np.einsum('ijklm,ijn->', a, b, optimize=True))
  503. # Issue #2027, was a problem in the contiguous 3-argument
  504. # inner loop implementation
  505. a = np.arange(1, 3)
  506. b = np.arange(1, 5).reshape(2, 2)
  507. c = np.arange(1, 9).reshape(4, 2)
  508. assert_equal(np.einsum('x,yx,zx->xzy', a, b, c),
  509. [[[1, 3], [3, 9], [5, 15], [7, 21]],
  510. [[8, 16], [16, 32], [24, 48], [32, 64]]])
  511. assert_equal(np.einsum('x,yx,zx->xzy', a, b, c, optimize=True),
  512. [[[1, 3], [3, 9], [5, 15], [7, 21]],
  513. [[8, 16], [16, 32], [24, 48], [32, 64]]])
  514. # Ensure explicitly setting out=None does not cause an error
  515. # see issue gh-15776 and issue gh-15256
  516. assert_equal(np.einsum('i,j', [1], [2], out=None), [[2]])
  517. def test_subscript_range(self):
  518. # Issue #7741, make sure that all letters of Latin alphabet (both uppercase & lowercase) can be used
  519. # when creating a subscript from arrays
  520. a = np.ones((2, 3))
  521. b = np.ones((3, 4))
  522. np.einsum(a, [0, 20], b, [20, 2], [0, 2], optimize=False)
  523. np.einsum(a, [0, 27], b, [27, 2], [0, 2], optimize=False)
  524. np.einsum(a, [0, 51], b, [51, 2], [0, 2], optimize=False)
  525. assert_raises(ValueError, lambda: np.einsum(a, [0, 52], b, [52, 2], [0, 2], optimize=False))
  526. assert_raises(ValueError, lambda: np.einsum(a, [-1, 5], b, [5, 2], [-1, 2], optimize=False))
  527. def test_einsum_broadcast(self):
  528. # Issue #2455 change in handling ellipsis
  529. # remove the 'middle broadcast' error
  530. # only use the 'RIGHT' iteration in prepare_op_axes
  531. # adds auto broadcast on left where it belongs
  532. # broadcast on right has to be explicit
  533. # We need to test the optimized parsing as well
  534. A = np.arange(2 * 3 * 4).reshape(2, 3, 4)
  535. B = np.arange(3)
  536. ref = np.einsum('ijk,j->ijk', A, B, optimize=False)
  537. for opt in [True, False]:
  538. assert_equal(np.einsum('ij...,j...->ij...', A, B, optimize=opt), ref)
  539. assert_equal(np.einsum('ij...,...j->ij...', A, B, optimize=opt), ref)
  540. assert_equal(np.einsum('ij...,j->ij...', A, B, optimize=opt), ref) # used to raise error
  541. A = np.arange(12).reshape((4, 3))
  542. B = np.arange(6).reshape((3, 2))
  543. ref = np.einsum('ik,kj->ij', A, B, optimize=False)
  544. for opt in [True, False]:
  545. assert_equal(np.einsum('ik...,k...->i...', A, B, optimize=opt), ref)
  546. assert_equal(np.einsum('ik...,...kj->i...j', A, B, optimize=opt), ref)
  547. assert_equal(np.einsum('...k,kj', A, B, optimize=opt), ref) # used to raise error
  548. assert_equal(np.einsum('ik,k...->i...', A, B, optimize=opt), ref) # used to raise error
  549. dims = [2, 3, 4, 5]
  550. a = np.arange(np.prod(dims)).reshape(dims)
  551. v = np.arange(dims[2])
  552. ref = np.einsum('ijkl,k->ijl', a, v, optimize=False)
  553. for opt in [True, False]:
  554. assert_equal(np.einsum('ijkl,k', a, v, optimize=opt), ref)
  555. assert_equal(np.einsum('...kl,k', a, v, optimize=opt), ref) # used to raise error
  556. assert_equal(np.einsum('...kl,k...', a, v, optimize=opt), ref)
  557. J, K, M = 160, 160, 120
  558. A = np.arange(J * K * M).reshape(1, 1, 1, J, K, M)
  559. B = np.arange(J * K * M * 3).reshape(J, K, M, 3)
  560. ref = np.einsum('...lmn,...lmno->...o', A, B, optimize=False)
  561. for opt in [True, False]:
  562. assert_equal(np.einsum('...lmn,lmno->...o', A, B,
  563. optimize=opt), ref) # used to raise error
  564. def test_einsum_fixedstridebug(self):
  565. # Issue #4485 obscure einsum bug
  566. # This case revealed a bug in nditer where it reported a stride
  567. # as 'fixed' (0) when it was in fact not fixed during processing
  568. # (0 or 4). The reason for the bug was that the check for a fixed
  569. # stride was using the information from the 2D inner loop reuse
  570. # to restrict the iteration dimensions it had to validate to be
  571. # the same, but that 2D inner loop reuse logic is only triggered
  572. # during the buffer copying step, and hence it was invalid to
  573. # rely on those values. The fix is to check all the dimensions
  574. # of the stride in question, which in the test case reveals that
  575. # the stride is not fixed.
  576. #
  577. # NOTE: This test is triggered by the fact that the default buffersize,
  578. # used by einsum, is 8192, and 3*2731 = 8193, is larger than that
  579. # and results in a mismatch between the buffering and the
  580. # striding for operand A.
  581. A = np.arange(2 * 3).reshape(2, 3).astype(np.float32)
  582. B = np.arange(2 * 3 * 2731).reshape(2, 3, 2731).astype(np.int16)
  583. es = np.einsum('cl, cpx->lpx', A, B)
  584. tp = np.tensordot(A, B, axes=(0, 0))
  585. assert_equal(es, tp)
  586. # The following is the original test case from the bug report,
  587. # made repeatable by changing random arrays to aranges.
  588. A = np.arange(3 * 3).reshape(3, 3).astype(np.float64)
  589. B = np.arange(3 * 3 * 64 * 64).reshape(3, 3, 64, 64).astype(np.float32)
  590. es = np.einsum('cl, cpxy->lpxy', A, B)
  591. tp = np.tensordot(A, B, axes=(0, 0))
  592. assert_equal(es, tp)
  593. def test_einsum_fixed_collapsingbug(self):
  594. # Issue #5147.
  595. # The bug only occurred when output argument of einssum was used.
  596. x = np.random.normal(0, 1, (5, 5, 5, 5))
  597. y1 = np.zeros((5, 5))
  598. np.einsum('aabb->ab', x, out=y1)
  599. idx = np.arange(5)
  600. y2 = x[idx[:, None], idx[:, None], idx, idx]
  601. assert_equal(y1, y2)
  602. def test_einsum_failed_on_p9_and_s390x(self):
  603. # Issues gh-14692 and gh-12689
  604. # Bug with signed vs unsigned char errored on power9 and s390x Linux
  605. tensor = np.random.random_sample((10, 10, 10, 10))
  606. x = np.einsum('ijij->', tensor)
  607. y = tensor.trace(axis1=0, axis2=2).trace()
  608. assert_allclose(x, y)
  609. def test_einsum_all_contig_non_contig_output(self):
  610. # Issue gh-5907, tests that the all contiguous special case
  611. # actually checks the contiguity of the output
  612. x = np.ones((5, 5))
  613. out = np.ones(10)[::2]
  614. correct_base = np.ones(10)
  615. correct_base[::2] = 5
  616. # Always worked (inner iteration is done with 0-stride):
  617. np.einsum('mi,mi,mi->m', x, x, x, out=out)
  618. assert_array_equal(out.base, correct_base)
  619. # Example 1:
  620. out = np.ones(10)[::2]
  621. np.einsum('im,im,im->m', x, x, x, out=out)
  622. assert_array_equal(out.base, correct_base)
  623. # Example 2, buffering causes x to be contiguous but
  624. # special cases do not catch the operation before:
  625. out = np.ones((2, 2, 2))[..., 0]
  626. correct_base = np.ones((2, 2, 2))
  627. correct_base[..., 0] = 2
  628. x = np.ones((2, 2), np.float32)
  629. np.einsum('ij,jk->ik', x, x, out=out)
  630. assert_array_equal(out.base, correct_base)
  631. @pytest.mark.parametrize("dtype",
  632. np.typecodes["AllFloat"] + np.typecodes["AllInteger"])
  633. def test_different_paths(self, dtype):
  634. # Test originally added to cover broken float16 path: gh-20305
  635. # Likely most are covered elsewhere, at least partially.
  636. dtype = np.dtype(dtype)
  637. # Simple test, designed to excersize most specialized code paths,
  638. # note the +0.5 for floats. This makes sure we use a float value
  639. # where the results must be exact.
  640. arr = (np.arange(7) + 0.5).astype(dtype)
  641. scalar = np.array(2, dtype=dtype)
  642. # contig -> scalar:
  643. res = np.einsum('i->', arr)
  644. assert res == arr.sum()
  645. # contig, contig -> contig:
  646. res = np.einsum('i,i->i', arr, arr)
  647. assert_array_equal(res, arr * arr)
  648. # noncontig, noncontig -> contig:
  649. res = np.einsum('i,i->i', arr.repeat(2)[::2], arr.repeat(2)[::2])
  650. assert_array_equal(res, arr * arr)
  651. # contig + contig -> scalar
  652. assert np.einsum('i,i->', arr, arr) == (arr * arr).sum()
  653. # contig + scalar -> contig (with out)
  654. out = np.ones(7, dtype=dtype)
  655. res = np.einsum('i,->i', arr, dtype.type(2), out=out)
  656. assert_array_equal(res, arr * dtype.type(2))
  657. # scalar + contig -> contig (with out)
  658. res = np.einsum(',i->i', scalar, arr)
  659. assert_array_equal(res, arr * dtype.type(2))
  660. # scalar + contig -> scalar
  661. res = np.einsum(',i->', scalar, arr)
  662. # Use einsum to compare to not have difference due to sum round-offs:
  663. assert res == np.einsum('i->', scalar * arr)
  664. # contig + scalar -> scalar
  665. res = np.einsum('i,->', arr, scalar)
  666. # Use einsum to compare to not have difference due to sum round-offs:
  667. assert res == np.einsum('i->', scalar * arr)
  668. # contig + contig + contig -> scalar
  669. arr = np.array([0.5, 0.5, 0.25, 4.5, 3.], dtype=dtype)
  670. res = np.einsum('i,i,i->', arr, arr, arr)
  671. assert_array_equal(res, (arr * arr * arr).sum())
  672. # four arrays:
  673. res = np.einsum('i,i,i,i->', arr, arr, arr, arr)
  674. assert_array_equal(res, (arr * arr * arr * arr).sum())
  675. def test_small_boolean_arrays(self):
  676. # See gh-5946.
  677. # Use array of True embedded in False.
  678. a = np.zeros((16, 1, 1), dtype=np.bool_)[:2]
  679. a[...] = True
  680. out = np.zeros((16, 1, 1), dtype=np.bool_)[:2]
  681. tgt = np.ones((2, 1, 1), dtype=np.bool_)
  682. res = np.einsum('...ij,...jk->...ik', a, a, out=out)
  683. assert_equal(res, tgt)
  684. def test_out_is_res(self):
  685. a = np.arange(9).reshape(3, 3)
  686. res = np.einsum('...ij,...jk->...ik', a, a, out=a)
  687. assert res is a
  688. def optimize_compare(self, subscripts, operands=None):
  689. # Tests all paths of the optimization function against
  690. # conventional einsum
  691. if operands is None:
  692. args = [subscripts]
  693. terms = subscripts.split('->')[0].split(',')
  694. for term in terms:
  695. dims = [global_size_dict[x] for x in term]
  696. args.append(np.random.rand(*dims))
  697. else:
  698. args = [subscripts] + operands
  699. noopt = np.einsum(*args, optimize=False)
  700. opt = np.einsum(*args, optimize='greedy')
  701. assert_almost_equal(opt, noopt)
  702. opt = np.einsum(*args, optimize='optimal')
  703. assert_almost_equal(opt, noopt)
  704. def test_hadamard_like_products(self):
  705. # Hadamard outer products
  706. self.optimize_compare('a,ab,abc->abc')
  707. self.optimize_compare('a,b,ab->ab')
  708. def test_index_transformations(self):
  709. # Simple index transformation cases
  710. self.optimize_compare('ea,fb,gc,hd,abcd->efgh')
  711. self.optimize_compare('ea,fb,abcd,gc,hd->efgh')
  712. self.optimize_compare('abcd,ea,fb,gc,hd->efgh')
  713. def test_complex(self):
  714. # Long test cases
  715. self.optimize_compare('acdf,jbje,gihb,hfac,gfac,gifabc,hfac')
  716. self.optimize_compare('acdf,jbje,gihb,hfac,gfac,gifabc,hfac')
  717. self.optimize_compare('cd,bdhe,aidb,hgca,gc,hgibcd,hgac')
  718. self.optimize_compare('abhe,hidj,jgba,hiab,gab')
  719. self.optimize_compare('bde,cdh,agdb,hica,ibd,hgicd,hiac')
  720. self.optimize_compare('chd,bde,agbc,hiad,hgc,hgi,hiad')
  721. self.optimize_compare('chd,bde,agbc,hiad,bdi,cgh,agdb')
  722. self.optimize_compare('bdhe,acad,hiab,agac,hibd')
  723. def test_collapse(self):
  724. # Inner products
  725. self.optimize_compare('ab,ab,c->')
  726. self.optimize_compare('ab,ab,c->c')
  727. self.optimize_compare('ab,ab,cd,cd->')
  728. self.optimize_compare('ab,ab,cd,cd->ac')
  729. self.optimize_compare('ab,ab,cd,cd->cd')
  730. self.optimize_compare('ab,ab,cd,cd,ef,ef->')
  731. def test_expand(self):
  732. # Outer products
  733. self.optimize_compare('ab,cd,ef->abcdef')
  734. self.optimize_compare('ab,cd,ef->acdf')
  735. self.optimize_compare('ab,cd,de->abcde')
  736. self.optimize_compare('ab,cd,de->be')
  737. self.optimize_compare('ab,bcd,cd->abcd')
  738. self.optimize_compare('ab,bcd,cd->abd')
  739. def test_edge_cases(self):
  740. # Difficult edge cases for optimization
  741. self.optimize_compare('eb,cb,fb->cef')
  742. self.optimize_compare('dd,fb,be,cdb->cef')
  743. self.optimize_compare('bca,cdb,dbf,afc->')
  744. self.optimize_compare('dcc,fce,ea,dbf->ab')
  745. self.optimize_compare('fdf,cdd,ccd,afe->ae')
  746. self.optimize_compare('abcd,ad')
  747. self.optimize_compare('ed,fcd,ff,bcf->be')
  748. self.optimize_compare('baa,dcf,af,cde->be')
  749. self.optimize_compare('bd,db,eac->ace')
  750. self.optimize_compare('fff,fae,bef,def->abd')
  751. self.optimize_compare('efc,dbc,acf,fd->abe')
  752. self.optimize_compare('ba,ac,da->bcd')
  753. def test_inner_product(self):
  754. # Inner products
  755. self.optimize_compare('ab,ab')
  756. self.optimize_compare('ab,ba')
  757. self.optimize_compare('abc,abc')
  758. self.optimize_compare('abc,bac')
  759. self.optimize_compare('abc,cba')
  760. def test_random_cases(self):
  761. # Randomly built test cases
  762. self.optimize_compare('aab,fa,df,ecc->bde')
  763. self.optimize_compare('ecb,fef,bad,ed->ac')
  764. self.optimize_compare('bcf,bbb,fbf,fc->')
  765. self.optimize_compare('bb,ff,be->e')
  766. self.optimize_compare('bcb,bb,fc,fff->')
  767. self.optimize_compare('fbb,dfd,fc,fc->')
  768. self.optimize_compare('afd,ba,cc,dc->bf')
  769. self.optimize_compare('adb,bc,fa,cfc->d')
  770. self.optimize_compare('bbd,bda,fc,db->acf')
  771. self.optimize_compare('dba,ead,cad->bce')
  772. self.optimize_compare('aef,fbc,dca->bde')
  773. def test_combined_views_mapping(self):
  774. # gh-10792
  775. a = np.arange(9).reshape(1, 1, 3, 1, 3)
  776. b = np.einsum('bbcdc->d', a)
  777. assert_equal(b, [12])
  778. def test_broadcasting_dot_cases(self):
  779. # Ensures broadcasting cases are not mistaken for GEMM
  780. a = np.random.rand(1, 5, 4)
  781. b = np.random.rand(4, 6)
  782. c = np.random.rand(5, 6)
  783. d = np.random.rand(10)
  784. self.optimize_compare('ijk,kl,jl', operands=[a, b, c])
  785. self.optimize_compare('ijk,kl,jl,i->i', operands=[a, b, c, d])
  786. e = np.random.rand(1, 1, 5, 4)
  787. f = np.random.rand(7, 7)
  788. self.optimize_compare('abjk,kl,jl', operands=[e, b, c])
  789. self.optimize_compare('abjk,kl,jl,ab->ab', operands=[e, b, c, f])
  790. # Edge case found in gh-11308
  791. g = np.arange(64).reshape(2, 4, 8)
  792. self.optimize_compare('obk,ijk->ioj', operands=[g, g])
  793. def test_output_order(self):
  794. # Ensure output order is respected for optimize cases, the below
  795. # conraction should yield a reshaped tensor view
  796. # gh-16415
  797. a = np.ones((2, 3, 5), order='F')
  798. b = np.ones((4, 3), order='F')
  799. for opt in [True, False]:
  800. tmp = np.einsum('...ft,mf->...mt', a, b, order='a', optimize=opt)
  801. assert_(tmp.flags.f_contiguous)
  802. tmp = np.einsum('...ft,mf->...mt', a, b, order='f', optimize=opt)
  803. assert_(tmp.flags.f_contiguous)
  804. tmp = np.einsum('...ft,mf->...mt', a, b, order='c', optimize=opt)
  805. assert_(tmp.flags.c_contiguous)
  806. tmp = np.einsum('...ft,mf->...mt', a, b, order='k', optimize=opt)
  807. assert_(tmp.flags.c_contiguous is False)
  808. assert_(tmp.flags.f_contiguous is False)
  809. tmp = np.einsum('...ft,mf->...mt', a, b, optimize=opt)
  810. assert_(tmp.flags.c_contiguous is False)
  811. assert_(tmp.flags.f_contiguous is False)
  812. c = np.ones((4, 3), order='C')
  813. for opt in [True, False]:
  814. tmp = np.einsum('...ft,mf->...mt', a, c, order='a', optimize=opt)
  815. assert_(tmp.flags.c_contiguous)
  816. d = np.ones((2, 3, 5), order='C')
  817. for opt in [True, False]:
  818. tmp = np.einsum('...ft,mf->...mt', d, c, order='a', optimize=opt)
  819. assert_(tmp.flags.c_contiguous)
  820. class TestEinsumPath:
  821. def build_operands(self, string, size_dict=global_size_dict):
  822. # Builds views based off initial operands
  823. operands = [string]
  824. terms = string.split('->')[0].split(',')
  825. for term in terms:
  826. dims = [size_dict[x] for x in term]
  827. operands.append(np.random.rand(*dims))
  828. return operands
  829. def assert_path_equal(self, comp, benchmark):
  830. # Checks if list of tuples are equivalent
  831. ret = (len(comp) == len(benchmark))
  832. assert_(ret)
  833. for pos in range(len(comp) - 1):
  834. ret &= isinstance(comp[pos + 1], tuple)
  835. ret &= (comp[pos + 1] == benchmark[pos + 1])
  836. assert_(ret)
  837. def test_memory_contraints(self):
  838. # Ensure memory constraints are satisfied
  839. outer_test = self.build_operands('a,b,c->abc')
  840. path, path_str = np.einsum_path(*outer_test, optimize=('greedy', 0))
  841. self.assert_path_equal(path, ['einsum_path', (0, 1, 2)])
  842. path, path_str = np.einsum_path(*outer_test, optimize=('optimal', 0))
  843. self.assert_path_equal(path, ['einsum_path', (0, 1, 2)])
  844. long_test = self.build_operands('acdf,jbje,gihb,hfac')
  845. path, path_str = np.einsum_path(*long_test, optimize=('greedy', 0))
  846. self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
  847. path, path_str = np.einsum_path(*long_test, optimize=('optimal', 0))
  848. self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
  849. def test_long_paths(self):
  850. # Long complex cases
  851. # Long test 1
  852. long_test1 = self.build_operands('acdf,jbje,gihb,hfac,gfac,gifabc,hfac')
  853. path, path_str = np.einsum_path(*long_test1, optimize='greedy')
  854. self.assert_path_equal(path, ['einsum_path',
  855. (3, 6), (3, 4), (2, 4), (2, 3), (0, 2), (0, 1)])
  856. path, path_str = np.einsum_path(*long_test1, optimize='optimal')
  857. self.assert_path_equal(path, ['einsum_path',
  858. (3, 6), (3, 4), (2, 4), (2, 3), (0, 2), (0, 1)])
  859. # Long test 2
  860. long_test2 = self.build_operands('chd,bde,agbc,hiad,bdi,cgh,agdb')
  861. path, path_str = np.einsum_path(*long_test2, optimize='greedy')
  862. self.assert_path_equal(path, ['einsum_path',
  863. (3, 4), (0, 3), (3, 4), (1, 3), (1, 2), (0, 1)])
  864. path, path_str = np.einsum_path(*long_test2, optimize='optimal')
  865. self.assert_path_equal(path, ['einsum_path',
  866. (0, 5), (1, 4), (3, 4), (1, 3), (1, 2), (0, 1)])
  867. def test_edge_paths(self):
  868. # Difficult edge cases
  869. # Edge test1
  870. edge_test1 = self.build_operands('eb,cb,fb->cef')
  871. path, path_str = np.einsum_path(*edge_test1, optimize='greedy')
  872. self.assert_path_equal(path, ['einsum_path', (0, 2), (0, 1)])
  873. path, path_str = np.einsum_path(*edge_test1, optimize='optimal')
  874. self.assert_path_equal(path, ['einsum_path', (0, 2), (0, 1)])
  875. # Edge test2
  876. edge_test2 = self.build_operands('dd,fb,be,cdb->cef')
  877. path, path_str = np.einsum_path(*edge_test2, optimize='greedy')
  878. self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 1), (0, 1)])
  879. path, path_str = np.einsum_path(*edge_test2, optimize='optimal')
  880. self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 1), (0, 1)])
  881. # Edge test3
  882. edge_test3 = self.build_operands('bca,cdb,dbf,afc->')
  883. path, path_str = np.einsum_path(*edge_test3, optimize='greedy')
  884. self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
  885. path, path_str = np.einsum_path(*edge_test3, optimize='optimal')
  886. self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
  887. # Edge test4
  888. edge_test4 = self.build_operands('dcc,fce,ea,dbf->ab')
  889. path, path_str = np.einsum_path(*edge_test4, optimize='greedy')
  890. self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 1), (0, 1)])
  891. path, path_str = np.einsum_path(*edge_test4, optimize='optimal')
  892. self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
  893. # Edge test5
  894. edge_test4 = self.build_operands('a,ac,ab,ad,cd,bd,bc->',
  895. size_dict={"a": 20, "b": 20, "c": 20, "d": 20})
  896. path, path_str = np.einsum_path(*edge_test4, optimize='greedy')
  897. self.assert_path_equal(path, ['einsum_path', (0, 1), (0, 1, 2, 3, 4, 5)])
  898. path, path_str = np.einsum_path(*edge_test4, optimize='optimal')
  899. self.assert_path_equal(path, ['einsum_path', (0, 1), (0, 1, 2, 3, 4, 5)])
  900. def test_path_type_input(self):
  901. # Test explicit path handling
  902. path_test = self.build_operands('dcc,fce,ea,dbf->ab')
  903. path, path_str = np.einsum_path(*path_test, optimize=False)
  904. self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
  905. path, path_str = np.einsum_path(*path_test, optimize=True)
  906. self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 1), (0, 1)])
  907. exp_path = ['einsum_path', (0, 2), (0, 2), (0, 1)]
  908. path, path_str = np.einsum_path(*path_test, optimize=exp_path)
  909. self.assert_path_equal(path, exp_path)
  910. # Double check einsum works on the input path
  911. noopt = np.einsum(*path_test, optimize=False)
  912. opt = np.einsum(*path_test, optimize=exp_path)
  913. assert_almost_equal(noopt, opt)
  914. def test_path_type_input_internal_trace(self):
  915. #gh-20962
  916. path_test = self.build_operands('cab,cdd->ab')
  917. exp_path = ['einsum_path', (1,), (0, 1)]
  918. path, path_str = np.einsum_path(*path_test, optimize=exp_path)
  919. self.assert_path_equal(path, exp_path)
  920. # Double check einsum works on the input path
  921. noopt = np.einsum(*path_test, optimize=False)
  922. opt = np.einsum(*path_test, optimize=exp_path)
  923. assert_almost_equal(noopt, opt)
  924. def test_path_type_input_invalid(self):
  925. path_test = self.build_operands('ab,bc,cd,de->ae')
  926. exp_path = ['einsum_path', (2, 3), (0, 1)]
  927. assert_raises(RuntimeError, np.einsum, *path_test, optimize=exp_path)
  928. assert_raises(
  929. RuntimeError, np.einsum_path, *path_test, optimize=exp_path)
  930. path_test = self.build_operands('a,a,a->a')
  931. exp_path = ['einsum_path', (1,), (0, 1)]
  932. assert_raises(RuntimeError, np.einsum, *path_test, optimize=exp_path)
  933. assert_raises(
  934. RuntimeError, np.einsum_path, *path_test, optimize=exp_path)
  935. def test_spaces(self):
  936. #gh-10794
  937. arr = np.array([[1]])
  938. for sp in itertools.product(['', ' '], repeat=4):
  939. # no error for any spacing
  940. np.einsum('{}...a{}->{}...a{}'.format(*sp), arr)
  941. def test_overlap():
  942. a = np.arange(9, dtype=int).reshape(3, 3)
  943. b = np.arange(9, dtype=int).reshape(3, 3)
  944. d = np.dot(a, b)
  945. # sanity check
  946. c = np.einsum('ij,jk->ik', a, b)
  947. assert_equal(c, d)
  948. #gh-10080, out overlaps one of the operands
  949. c = np.einsum('ij,jk->ik', a, b, out=b)
  950. assert_equal(c, d)