Given an n-by-n matrix A, where each row of A is a permutation of [n], e.g.,
import torch
n = 100
AA = torch.rand(n, n)
A = torch.argsort(AA, dim=1)
Also given another n-by-n matrix P, we want to construct a 3D tensor Q s.t.
Q[i, j, k] = P[A[i, j], k]
Is there any efficient way in pytorch?
I am aware of torch.gather but it seems hard to be directly applied here.
You can directly use:
Q = P[A]
Why not simply use A as an index:
Q = P[A, :]
I'm trying to vectorize a loop with NumPy but I'm stuck
I have a matrix A of shape (NN,NN) I define the A-dot product by
def scalA(u,v):
return v.T # A # u
Then I have two matrices B and C (B has a shape (N,NN) and C has a shape (K,NN) the loop I'm trying to vectorize is
res = np.zeros((N,K))
for n in range(N):
for k in range(K):
res[n,k] = scalA(B[n,:], C[k,:])
I found during my research functions like np.tensordot or np.einsum, but I haven't really understood how they work, and (if I have well understood) tensordot will compute the canonical dot product (that would correspond to A = np.eye(NN) in my case).
Thanks !
np.einsum('ni,ji,kj->nk', B,A,C)
I think this works. I wrote it 'by eye' without testing.
Probably you're looking for this:
def scalA(u,v):
return u # A # v.T
If shape of A is (NN,NN), shape of B is (N,NN), and shape of C is (K,NN), the result of scalA(B,C) has shape (N,K)
If shape of A is (NN,NN), shape of B is (NN,), and shape of C is (NN,), the result of scalA(B,C) is just a scalar.
However, if you're expecting B and C to have even higher dimensionality (greater than 2), this may need further tweaking. (I could not tell from your question whether that's the case)
I have two matrices A, B, NxKxD dimensions and I want get matrix C, NxKxDxD dimensions, where C[n, k] = A[n, k] x B[n, k].T (here "x" means product of matrices of dimensions Dx1 and 1xD, so the result must be DxD dimensional), so now my code looking like this (here A = B = X):
def square(X):
out = np.zeros((N, K, D, D))
for n in range(N):
for k in range(K):
out[n, k] = np.dot(X[n, k, :, np.newaxis], X[n, k, np.newaxis, :])
return out
It may be slow for big N and K because of python's for cycle. Is there some way to make this multiplication in one numpy function?
It seems you are not using np.dot for sum-reduction, but just for expansion that results in broadcasting. So, you can simply extend the array to have one more dimension with the use of np.newaxis/None and let the implicit broadcasting help out.
Thus, an implementation would be -
X[...,None]*X[...,None,:]
More info on broadcasting specifically how to add new axes could be found in this other post.
I'm working with numpy arrays of shape (N,), (N,3) and (N,3,3) which represent sequences of scalars, vectors and matrices in 3D space. I have implemented pointwise dot product, matrix multiplication, and matrix/vector multiplication as follows:
def dot_product(v, w):
return np.einsum('ij, ij -> i', v, w)
def matrix_vector_product(M, v):
return np.einsum('ijk, ik -> ij', M, v)
def matrix_matrix_product(A, B):
return np.einsum('ijk, ikl -> ijl', A, B)
As you can see I use einsum for lack of a better solution. To my surprise I was not able to use np.dot... which seems not suitable for this need. Is there a more numpythonic way to implement these function?
In particular it would be nice if the functions could work also on the shapes (3,) and (3,3) by broadcasting the first missing axis. I think I need ellipsis, but I don't quite understand how to achieve the result.
These operations cannot be reshaped into general BLAS calls and looping BLAS calls would be quite slow for arrays of this size. As such, einsum is likely optimal for this kind of operation.
Your functions can be generalized with ellipses as follows:
def dot_product(v, w):
return np.einsum('...j,...j->...', v, w)
def matrix_vector_product(M, v):
return np.einsum('...jk,...k->...j', M, v)
def matrix_matrix_product(A, B):
return np.einsum('...jk,...kl->...jl', A, B)
Just as working notes, these 3 calculations can also be written as:
np.einsum(A,[0,1,2],B,[0,2,3],[0,1,3])
np.einsum(M,[0,1,2],v,[0,2],[0,1])
np.einsum(w,[0,1],v,[0,1],[0])
Or with Ophion's generalization
np.einsum(A,[Ellipsis,1,2], B, ...)
It shouldn't be hard to generate the [0,1,..] lists based on the dimensions of the inputs arrays.
By focusing on generalizing the einsum expressions, I missed the fact that what you are trying to reproduce is N small dot products.
np.array([np.dot(i,j) for i,j in zip(a,b)])
It's worth keeping mind that np.dot uses fast compiled code, and focuses on calculations where the arrays are large. Where as your problem is one of calculating many small dot products.
And without extra arguments that define axes, np.dot performs just 2 of the possible combinations, ones which can be expressed as:
np.einsum('i,i', v1, v2)
np.einsum('...ij,...jk->...ik', m1, m2)
An operator version of dot would face the same limitation - no extra parameters to specify how the axes are to be combined.
It may also be instructive to note what tensordot does to generalize dot:
def tensordot(a, b, axes=2):
....
newshape_a = (-1, N2)
...
newshape_b = (N2, -1)
....
at = a.transpose(newaxes_a).reshape(newshape_a)
bt = b.transpose(newaxes_b).reshape(newshape_b)
res = dot(at, bt)
return res.reshape(olda + oldb)
It can perform a dot with summation over several axes. But after the transposing and reshaping is done, the calculation becomes the standard dot with 2d arrays.
This could have been flagged as a duplicate issue. People have asking about doing multiple dot products for some time.
Matrix vector multiplication along array axes
suggests using numpy.core.umath_tests.matrix_multiply
https://stackoverflow.com/a/24174347/901925 equates:
matrix_multiply(matrices, vectors[..., None])
np.einsum('ijk,ik->ij', matrices, vectors)
The C documentation for matrix_multiply notes:
* This implements the function
* out[k, m, p] = sum_n { in1[k, m, n] * in2[k, n, p] }.
inner1d from the same directory does the same same for (N,n) vectors
inner1d(vector, vector)
np.einsum('ij,ij->i', vector, vector)
# out[n] = sum_i { in1[n, i] * in2[n, i] }
Both are UFunc, and can handle broadcasting on the right most dimensions. In numpy/core/test/test_ufunc.py these functions are used to exercise the UFunc mechanism.
matrix_multiply(np.ones((4,5,6,2,3)),np.ones((3,2)))
https://stackoverflow.com/a/16704079/901925 adds that this kind of calculation can be done with * and sum, eg
(w*v).sum(-1)
(M*v[...,None]).sum(-1)
(A*B.swapaxes(...)).sum(-1)
On further testing, I think inner1d and matrix_multiply match your dot and matrix-matrix product cases, and the matrix-vector case if you add the [...,None]. Looks like they are 2x faster than the einsum versions (on my machine and test arrays).
https://github.com/numpy/numpy/blob/master/doc/neps/return-of-revenge-of-matmul-pep.rst
is the discussion of the # infix operator on numpy. I think the numpy developers are less enthused about this PEP than the Python ones.
This is probably obvious on reflection, but it's not clear to me right now.
For a pair of numpy arrays of shapes (K, N, M) and (K, M, N) denoted by a and b respectively, is there a way to compute the following as a single vectorized operation:
import numpy as np
K = 5
N = 2
M = 3
a = np.random.randn(K, N, M)
b = np.random.randn(K, M, N)
output = np.empty((K, N, N))
for each_a, each_b, each_out in zip(a, b, output):
each_out[:] = each_a.dot(each_b)
A simple a.dot(b) returns the dot product for every pair of the first axis (so it returns an array of shape (K, N, K, N).
edit: fleshed out the code a bit for those that couldn't understand the question.
I answered a similar question a while back: Element-wise matrix multiplication in NumPy .
I think what you're looking for is:
output = np.einsum('ijk,ikl->ijl', a, b)
Good luck!