由于没有要对齐的轴,我们可以简单地使用tensordot
它让不参与减和的轴通过额外的“展开” rollaxis
,就像这样 -
np.rollaxis(np.tensordot(a,b,axes=(1,0)),a.ndim-1,1)
如果你想使用einsum
,我们可以将它们重新整形3D
,使它们的最后一个轴是合并的轴(第三个轴向前合并为一个),然后继续进行einsum
最后重新整形,使其ndim-1
形状在输出中展开,像这样的东西 -
shp_a = a.shape
shp_b = b.shape
shp_a[:1] + shp_a[2:]
out_shp = shp_a[:1] + (shp_b[1],) + shp_a[2:] + shp_b[2:]
a3D = a.reshape(shp_a[:2]+(-1,))
b3D = b.reshape(shp_b[:2]+(-1,))
out = np.einsum('ijk,jlm->ilkm',a3D,b3D).reshape(out_shp)
我们还可以自己生成相应的 einsum 字符串表示法,从而跳过所有数组操作,从而专注于字符串操作本身以获得类似的结果 -
import string
def einsum_spreadout(a,b,a_axes,b_axes,a_spread_axis,b_spread_axis):
from numpy.core import numerictypes as nt
if isinstance(a_axes, (int, nt.integer)):
a_axes = (a_axes,)
if isinstance(b_axes, (int, nt.integer)):
b_axes = (b_axes,)
s = string.ascii_letters
a_str = s[:a.ndim]
b_str = s[a.ndim:a.ndim+b.ndim]
b_str_ar = np.frombuffer(b_str,dtype='S1').copy()
for (i,j) in zip(a_axes,b_axes):
b_str_ar[j] = a_str[i]
b_str = ''.join(b_str_ar)
out_str = a_str[:a_spread_axis] + b_str[:b_spread_axis]
out_str += a_str[a_spread_axis:] + b_str[b_spread_axis:]
out_str_ar = np.frombuffer(out_str,dtype='S1').copy()
out_str = ''.join(out_str_ar[~np.isin(out_str_ar,np.take(b_str_ar,b_axes))])
einsum_str = a_str+','+b_str+'->'+out_str
return np.einsum(einsum_str,a,b)
很少有示例案例运行以展示其用法 -
>>> a = np.random.rand(3,4,6,7,8)
>>> b = np.random.rand(4,5,9,10)
>>> einsum_spreadout(a,b,a_axes=1,b_axes=0,a_spread_axis=2,b_spread_axis=2).shape
(3, 5, 6, 7, 8, 9, 10)
>>> b = np.random.rand(4,5,6,10)
>>> einsum_spreadout(a,b,a_axes=(1,2),b_axes=(0,2),a_spread_axis=2,b_spread_axis=2).shape
(3, 5, 7, 8, 10)
>>> einsum_spreadout(a,b,a_axes=(1,2),b_axes=(0,2),a_spread_axis=4,b_spread_axis=4).shape
(3, 7, 5, 10, 8)