We Keep Coding

Matrix multiplication while subsetting elements from matrices and storing in a new matrix

I am attempting a numpy.matmul call using as variables
Matrix A of dimensions (p, t, q)
Matrix B of dimensions (r, t).
A categories vector of shape r and p categories, used to take slices of B and define the index of A do use.
The multiplications are done iteratively using the indices of each category. For each category p_i, I extract from A a submatrix (t, q). Then, I multiply those with a subset of B (x, t), where x is a mask defined by r == p_i. Finally, the matrix multiplication of (x, t) and (t, q) produces the output (x, q) which is stored at S[x].
I have noted that I cannot figure out a non-iterative version of this algorithm. The first snippet describes an iterative solution. The second one is an attempt at what I would wish to get, where everything is calculated as a single-step and would be presumably faster. However, it is incorrect because matrix A has three dimensions instead of two. Maybe there is no way to do this in NumPy with a single call, and in general, looking for advice/ideas to try out.
Thanks!
import numpy as np
p, q, r, t = 2, 9, 512, 4
# data initialization (random)
np.random.seed(500)
S = np.random.rand(r, q)
A = np.random.randint(0, 3, size=(p, t, q))
B = np.random.rand(r, t)
categories = np.random.randint(0, p, r)
print('iterative') # iterative
for i in range(p):
# print(i)
a = A[i, :, :]
mask = categories == i
b = B[mask]
print(b.shape, a.shape, S[mask].shape,
np.matmul(b, a).shape)
S[mask] = np.matmul(b, a)
print(S.shape)
a simple way to write it down
S = np.random.rand(r, q)
print(A[:p,:,:].shape)
result = np.matmul(B, A[:p,:,:])
# iterative assignment
i = 0
S[categories == i] = result[i, categories == i, :]
i = 1
S[categories == i] = result[i, categories == i, :]
The next snippet will produce an error during the multiplication step.
# attempt to multiply once, indexing all categories only once (not possible)
np.random.seed(500)
S = np.random.rand(r, q)
# attempt to use the categories vector
a = A[categories, :, :]
b = B[categories]
# due to the shapes of the arrays, this multiplication is not possible
print('\nsingle step (error due to shapes of the matrix a')
print(b.shape, a.shape, S[categories].shape)
S[categories] = np.matmul(b, a)
print(scores.shape)
iterative
(250, 4) (4, 9) (250, 9) (250, 9)
(262, 4) (4, 9) (262, 9) (262, 9)
(512, 9)
single step (error due to shapes of the 2nd matrix a).
(512, 4) (512, 4, 9) (512, 9)

In [63]: (np.ones((512,4))#np.ones((512,4,9))).shape
Out[63]: (512, 512, 9)
This because the first array is broadcasted to (1,512,4). I think you want instead to do:
In [64]: (np.ones((512,1,4))#np.ones((512,4,9))).shape
Out[64]: (512, 1, 9)
Then remove the middle dimension to get a (512,9).
Another way:
In [72]: np.einsum('ij,ijk->ik', np.ones((512,4)), np.ones((512,4,9))).shape
Out[72]: (512, 9)

To remove the loop altogether, you can try this
bigmask = np.arange(p)[:, np.newaxis] == categories
C = np.matmul(B, A)
res = C[np.broadcast_to(bigmask[..., np.newaxis], C.shape)].reshape(r, q)
# `res` has the same rows as the iterative `S` but in the wrong order
# so we need to reorder the rows
sort_index = np.argsort(np.broadcast_to(np.arange(r), bigmask.shape)[bigmask])
assert np.allclose(S, res[sort_index])
Though I'm not sure it's much faster than the iterative version.

Vectorizing a function that takes multidimensional input over a multidimensional array in JAX

I have been trying to vectorize a function that takes two 2D arrays and return a
2D array of the same shape, so that I can apply it element wise to two 4D arrays.
Here is an example:
import jax.numpy as jnp
from jax import map
M = jnp.arange(0, 400).reshape((10, 10, 2, 2))
N = jnp.arange(400, 800).reshape((10, 10, 2, 2))
def func(A, B):
return jnp.dot(A, B)
R = vmap(func, in_axes=(1, 1))(M, N)
print(R.shape) #(10, 10, 2, 10, 2)
func is the function that I would like to apply to the 2x2 matrices that are contained in M and N. I expected the result to be of the shape (10, 10, 2, 2) because I thought that vmap applies the function to each subarray of the specified axis. Clearly I am not correctly understanding how it works. I appreciate any help! Thanks!

vmap will map over a single axis at a time. Since you want to map over two axes in each array, you should use two vmap calls:
R = vmap(vmap(func))(M, N)
print(R.shape) #(10, 10, 2, 2)
The reason a single vmap returns the shape it does is because you're mapping over a single axis of each array, so that the inputs to func effectively have shapes (10, 2, 2) and (10, 2, 2). If you call jnp.dot on arrays of these shapes, it will return an array of shape (10, 2, 10, 2) (the jnp.dot documentation makes clear why this is); adding the leading mapped dimension of size (10,) results in the shape you saw originally.

You could use jax.numpy.einsum too. This would allow you to cooperate with changing dimensions of M and N. Here is an example of 5 dimensional M and N:
import jax.numpy as jnp
from jax import vmap
from string import ascii_lowercase as letters
M = jnp.arange(0, 4000).reshape((10, 10, 10, 2, 2))
N = jnp.arange(4000, 8000).reshape((10, 10, 10, 2, 2))
C = letters[:(M.ndim - 2)] # You will map over first M.ndim-2 dimensions
subscripts = C+'XY,'+C+'YZ->'+C+'XZ' # Here I create the einsum script for the product
R1 = jnp.einsum(subscripts, M, N)
You could do the same thing with vmap. But you need three vmap calls this time:
def func(A, B):
return jnp.dot(A, B)
R2 = vmap(vmap(vmap(func)))(M, N)
assert jnp.allclose(R1, R2)

Most efficient way to perform large dot/tensor dot products while only keeping diagonal entries [duplicate]

This question already has answers here:
Matrix multiplication for multidimensional matrix (/array) - how to avoid loop?
(3 answers)
Closed 3 years ago.
I'm trying to figure out a way to use numpy to perform the following algebra in the most time-efficient way possible:
Given a 3D matrix/tensor, A, with shape (n, m, p) and a 2D matrix/tensor, B, with shape (n, p), calculate C_ij = sum_over_k (A_ijk * B_ik), where the resulting matrix C would have dimension (n, m).
I've tried two ways to do this. One is to loop through the first dimension, and calculate a regular dot product each time.
The other method is to use np.tensordot(A, B.T) to calculate a result with shape (n, m, n), and then take the diagonal elements along 1st and 3rd dimension. Both methods are shown below.
First method:
C = np.zeros((n,m))
for i in range(n):
C[i] = np.dot(A[i], B[i])
Second method:
C = np.diagonal(np.tensordot(A, B.T, axes = 1), axis1=0, axis2=2).T
However, because n is a very large number, the loop over n in the first method is costing a lot of time. The second method calculates too many unnecessary entries to obtain that huge (n, m, n)matrix, and is also costing too much time, I'm wondering if there's any efficient way to do this?

Define 2 arrays:
In [168]: A = np.arange(2*3*4).reshape(2,3,4); B = np.arange(2*4).reshape(2,4)
Your iterative approach:
In [169]: [np.dot(a,b) for a,b in zip(A,B)]
Out[169]: [array([14, 38, 62]), array([302, 390, 478])]
The einsum practically writes itself from your C_ij = sum_over_k (A_ijk * B_ik):
In [170]: np.einsum('ijk,ik->ij', A, B)
Out[170]:
array([[ 14, 38, 62],
[302, 390, 478]])
#, matmul, was added to perform batch dot products; here the i dimension is the batch one. Since it uses the last of A and 2nd to the last of B for the dot summation, we have to temporarily expand B to (2,4,1):
In [171]: A#B[...,None]
Out[171]:
array([[[ 14],
[ 38],
[ 62]],
[[302],
[390],
[478]]])
In [172]: (A#B[...,None])[...,0]
Out[172]:
array([[ 14, 38, 62],
[302, 390, 478]])
Typically matmul is fastest, since it passes the task to BLAS like code.

here is my implementation:
B = np.expand_dims(B, axis=1)
E = A * B
E = np.sum(E, axis=-1)
Check :
import numpy as np
n, m, p = 2, 2, 2
np.random.seed(0)
A = np.random.randint(1, 10, (n, m, p))
B = np.random.randint(1, 10, (n, p))
C = np.diagonal(np.tensordot(A, B.T, axes = 1), axis1=0, axis2=2).T
# from here is my implementation
B = np.expand_dims(B, axis=1)
E = A * B
E = np.sum(E, axis=-1)
print(np.array_equal(C, E))
True
use the np.expand_dims() to add a new dimension.
And use the broadcast multiply. Finally, sum along the third dimension.
Thanks check code from user3483203

multiplication of 3-dimensional matrix in numpy

I think I asked the wrong question yesterday. What I actually want is to mutiply two 2x2xN matrices A and B, so that
C[:,:,i] = dot(A[:,:,i], B[:,:,i])
For example, if I have a matrix
A = np.arange(12).reshape(2, 2, 3)
How can I get C = A x A with the definition described above? Is there a built-in function to do this?
Also, if I multiply A (shape 2x2xN) with B (shape 2x2x1, instead of N), I want to get
C[:,:,i] = dot(A[:,:,i], B[:,:,1])

Try using numpy.einsum, it has a little bit of a learning curve but it should give you what you want. Here is an example to get you started.
import numpy as np
A = np.random.random((2, 2, 3))
B = np.random.random((2, 2, 3))
C1 = np.empty((2, 2, 3))
for i in range(3):
C1[:, :, i] = np.dot(A[:, :, i], B[:, :, i])
C2 = np.einsum('ijn,jkn->ikn', A, B)
np.allclose(C1, C2)

Get the 1-norm of a vector in Python

How can I calculate the 1-norm of the difference of two vectors, ||a - b||_1 = sum(|a_i - b_i|) in Python?
a = [1,2,3,4]
b = [2,3,4,5]
||a - b||_1 = 4

Python has powerful built-in types, but Python lists are not mathematical vectors or matrices. You could do this with lists, but it will likely be cumbersome for anything more than trivial operations.
If you find yourself needing vector or matrix arithmetic often, the standard in the field is NumPy, which probably already comes packaged for your operating system the way Python also was.
I share the confusion of others about exactly what it is you're trying to do, but perhaps the numpy.linalg.norm function will help:
>>> import numpy
>>> a = numpy.array([1, 2, 3, 4])
>>> b = numpy.array([2, 3, 4, 5])
>>> numpy.linalg.norm((a - b), ord=1)
4
To show how that's working under the covers:
>>> a
array([1, 2, 3, 4])
>>> b
array([2, 3, 4, 5])
>>> (a - b)
array([-1, -1, -1, -1])
>>> numpy.linalg.norm((a - b))
2.0
>>> numpy.linalg.norm((a - b), ord=1)
4

In NumPy, for two vectors a and b, this is just
numpy.linalg.norm(a - b, ord=1)

You appear to be asking for the sum of the differences between the paired components of the two arrays:
>>> A=[1,2,3,4]
>>> B=[2,3,4,5]
>>> sum(abs(a - b) for a, b in zip(A, B))
4

It is not clear what exactly is required here, but here is my guess
a=[1,2,3,4]
b=[2,3,4,5]
def a_b(a,b):
return sum(map(lambda a:abs(a[0]-a[1]), zip(a,b)))
print a_b(a,b)

Using Numpy you can calculate any norm between two vectors using the linear algebra package.
import numpy as np
a = np.array([[2,3,4])
b = np.array([0,-1,7])
# L1 Norm
np.linalg.norm(a-b, ord=1)
# L2 Norm
np.linalg.norm(a-b, ord=2)
# L3 Norm
np.linalg.norm(a-b, ord=3)
# Ln Norm
np.linalg.norm(a-b, ord=n)
Example: