I have two arrays. One is n
by p
and the other is d
by p
by r
. I would like my output to be d
by n
by r
, which I can achieve easily as I construct the tensor B
below. However, I would like to do this without that loop.
import numpy
X = numpy.array([[1,2,3],[3,4,5],[5,6,7],[7,8,9]]) # n x p
betas = numpy.array([[[1,2],[1,2],[1,2]], [[5,6],[5,6],[5,6]]]) # d x p x r
print X.shape
print betas.shape
B = numpy.zeros((betas.shape[0],X.shape[0],betas.shape[2]))
print B.shape
for i in range(B.shape[0]):
B[i,:,:] = numpy.dot(X, betas[i])
print "B",B
C = numpy.tensordot(X, betas, axes=([1],[0]))
print C.shape
I have tried in various ways to get C
to match B
, but so far I have been unsuccessful. Is there a way that does not involve a call to reshape
?
Since the dot
rule is 'last of A with 2nd to the last of B', you can do X.dot(betas)
and get a (n,d,r) array (this sums on the shared p
dimension). Then you just need a transpose to get (d,n,r)
In [200]: X.dot(betas).transpose(1,0,2)
Out[200]:
array([[[ 6, 12],
[ 12, 24],
[ 18, 36],
[ 24, 48]],
[[ 30, 36],
[ 60, 72],
[ 90, 108],
[120, 144]]])
We can also write the einsum
version directly from the dimensions specification:
np.einsum('np,dpr->dnr', X,betas)
So does matmul
(this does dot
on the last 2 axes, while d
comes along for the ride).
X@betas
- If either argument is ND, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly.
We can use np.tensordot
and then need to permutes axes -
B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)
# Or np.tensordot(X, betas, axes=(1,1)).swapaxes(0,1)
Here is another approach using numpy.dot()
, which also returns a view as you requested, and most importantly more than 4x faster than tensordot approach, particularly for small sized arrays. But, np.tensordot
is way faster than plain np.dot()
for reasonably larger arrays. See timings below.
In [108]: X.shape
Out[108]: (4, 3)
In [109]: betas.shape
Out[109]: (2, 3, 2)
# use `np.dot` and roll the second axis to first position
In [110]: dot_prod = np.rollaxis(np.dot(X, betas), 1)
In [111]: dot_prod.shape
Out[111]: (2, 4, 2)
# @Divakar's approach
In [113]: B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)
# sanity check :)
In [115]: np.all(np.equal(dot_prod, B))
Out[115]: True
Now, the performance of both approaches:
np.dot()
is 4x faster than np.tensordot()
# @Divakar's approach
In [117]: %timeit B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)
10.6 µs ± 2.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
# @hpaulj's approach
In [151]: %timeit esum_dot = np.einsum('np, dpr -> dnr', X, betas)
4.16 µs ± 235 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# proposed approach: more than 4x faster!!
In [118]: %timeit dot_prod = np.rollaxis(np.dot(X, betas), 1)
2.47 µs ± 11.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
np.tensordot()
is much faster than np.dot()
In [129]: X = np.random.randint(1, 10, (600, 500))
In [130]: betas = np.random.randint(1, 7, (300, 500, 300))
In [131]: %timeit B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)
18.2 s ± 2.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [132]: %timeit dot_prod = np.rollaxis(np.dot(X, betas), 1)
52.8 s ± 14.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.