test_einsum.py 46 KB

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