maxframe.tensor.tensordot#

maxframe.tensor.tensordot(a, b, axes=2, sparse=None)[源代码]#

计算沿指定轴的张量点积,适用于维度大于等于1的张量。

给定两个张量(维度大于等于一的数组),ab,以及一个包含两个 array_like 对象的 array_like 对象 (a_axes, b_axes),对 ab 的元素(分量)在由 a_axesb_axes 指定的轴上求和。第三个参数可以是一个非负整数标量 N;如果是这样,则对 a 的最后 N 个维度和 b 的前 N 个维度求和。

参数:
  • a (array_like, len(shape) >= 1) -- 要进行“点积”的张量。

  • b (array_like, len(shape) >= 1) -- 要进行“点积”的张量。

  • axes (int or (2,) array_like) --

    • integer_like 如果是一个整数 N,则依次对 a 的最后 N 个轴和 b 的前 N 个轴求和。相应轴的大小必须匹配。

    • (2,) array_like 或者,一个包含要被求和的轴的列表,第一个序列应用于 a,第二个序列应用于 b。这两个 array_like 元素的长度必须相同。

参见

dot, einsum

备注

三种常见的使用场景是:

  • axes = 0:张量积 \(a\otimes b\)

  • axes = 1:张量点积 \(a\cdot b\)

  • axes = 2:(默认)张量双收缩 \(a:b\)

axes 是 integer_like 时,计算顺序是:首先 a 的第 -N 个轴和 b 的第 0 个轴,最后是 a 的第 -1 个轴和 b 的第 N 个轴。

当有多个轴需要求和时,且这些轴不是 a`(`b)的最后(前)轴时,参数 axes 应该包含两个相同长度的序列,第一个序列中的第一个轴和第二个序列中的第一个轴对应,第二个轴对应第二个,依此类推。

示例

>>> import maxframe.tensor as mt

一个“传统”的例子:

>>> a = mt.arange(60.).reshape(3,4,5)
>>> b = mt.arange(24.).reshape(4,3,2)
>>> c = mt.tensordot(a,b, axes=([1,0],[0,1]))
>>> c.shape
(5, 2)
>>> r = c.execute()
>>> r
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])
>>> # A slower but equivalent way of computing the same...
>>> ra = np.arange(60.).reshape(3,4,5)
>>> rb = np.arange(24.).reshape(4,3,2)
>>> d = np.zeros((5,2))
>>> for i in range(5):
...   for j in range(2):
...     for k in range(3):
...       for n in range(4):
...         d[i,j] += ra[k,n,i] * rb[n,k,j]
>>> r == d
array([[ True,  True],
       [ True,  True],
       [ True,  True],
       [ True,  True],
       [ True,  True]], dtype=bool)

一个利用 + 和 * 重载的扩展示例:

>>> a = mt.array(range(1, 9))
>>> a.shape = (2, 2, 2)
>>> A = mt.array(('a', 'b', 'c', 'd'), dtype=object)
>>> A.shape = (2, 2)
>>> a.execute(); A.execute()
array([[[1, 2],
        [3, 4]],
       [[5, 6],
        [7, 8]]])
array([[a, b],
       [c, d]], dtype=object)
>>> mt.tensordot(a, A).execute() # third argument default is 2 for double-contraction
array([abbcccdddd, aaaaabbbbbbcccccccdddddddd], dtype=object)
>>> mt.tensordot(a, A, 1).execute()
array([[[acc, bdd],
        [aaacccc, bbbdddd]],
       [[aaaaacccccc, bbbbbdddddd],
        [aaaaaaacccccccc, bbbbbbbdddddddd]]], dtype=object)
>>> mt.tensordot(a, A, 0).execute() # tensor product (result too long to incl.)
array([[[[[a, b],
          [c, d]],
          ...
>>> mt.tensordot(a, A, (0, 1)).execute()
array([[[abbbbb, cddddd],
        [aabbbbbb, ccdddddd]],
       [[aaabbbbbbb, cccddddddd],
        [aaaabbbbbbbb, ccccdddddddd]]], dtype=object)
>>> mt.tensordot(a, A, (2, 1)).execute()
array([[[abb, cdd],
        [aaabbbb, cccdddd]],
       [[aaaaabbbbbb, cccccdddddd],
        [aaaaaaabbbbbbbb, cccccccdddddddd]]], dtype=object)
>>> mt.tensordot(a, A, ((0, 1), (0, 1))).execute()
array([abbbcccccddddddd, aabbbbccccccdddddddd], dtype=object)
>>> mt.tensordot(a, A, ((2, 1), (1, 0))).execute()
array([acccbbdddd, aaaaacccccccbbbbbbdddddddd], dtype=object)