首页 > 解决方案 > numpy - 张量乘积

问题描述

我有一个 4 x 4 矩阵

import numpy as np
c = np.random.rand((4,4))

我想创建一个 100 x 4 x 4 x 100 张量,这样当第一个和最后一个索引相等时,我得到我的矩阵,否则我得到零。

我可以循环执行此操作

Z = np.zeros((100, 4, 4, 100))
for i in range(100):
    Z[i, :, :, i] = c

有一个更好的方法吗?我尝试查看 np.tensordot 和 np.einsum 但无法弄清楚。

谢谢, 萨希尔

标签: pythonnumpytensor

解决方案


使用advanced-indexing-

n = 100
Zout = np.zeros((n, 4, 4, n))
I = np.arange(n)
Zout[I,:,:,I] = c

eye-masking-

n = 100
mask = np.eye(n, dtype=bool)
Zout = np.zeros((n, 4, 4, n))
Zout.transpose(0,3,1,2)[mask] = c

计时 -

In [72]: c = np.random.rand(4,4)

In [73]: %%timeit
    ...: n = 100
    ...: Zout = np.zeros((n, 4, 4, n))
    ...: I = np.arange(n)
    ...: Zout[I,:,:,I] = c
10000 loops, best of 3: 47.5 µs per loop

In [74]: %%timeit
    ...: n = 100
    ...: mask = np.eye(n, dtype=bool)
    ...: Zout = np.zeros((n, 4, 4, n))
    ...: Zout.transpose(0,3,1,2)[mask] = c
10000 loops, best of 3: 73.1 µs per loop

推荐阅读