CODEWITHSUNDEEP
pythontensorflowmatrix-multiplication
bemma
bemma

Reputation: 95

3 dimensional matrix multiplication in tensorflow

For a 2-dimensional matrix A of size (N, K) with each element 'a', we can get a matrix B of size (N, K, N) with each element 'b' such that b[i, k, j] = a[i, k]*a[j,k] by the operation B = tf.expand_dims(A, -1)* tf.transpose(A).

Now with a matrix of 3-dimensional matrix A of size (M, N, K) with each element 'a', is there a way to compute 4-dimensional matrix B of size (M, N, K, N) with each element 'b' such that b[m, i, k, j] = a[m, i, k]*a[m, j, k]?

Upvotes: 3

Views: 520

Answers (2)

Ehsan
Ehsan

Reputation: 12407

Try einsum:

B = np.einsum('mik,mjk->mikj', A, A)

You can use (tf.einsum) if you are using tensors.

Upvotes: 2

Poe Dator
Poe Dator

Reputation: 4893

Bemma, This solution should work: Expand N dimension, multiply, transpose result.

M, N, K = 2,3,4  # insert your dimensions here
A = tf.constant(np.random.randint(1, 100, size=[M,N,K]))  # generate A
B = tf.expand_dims(A, 1)* tf.expand_dims(A, 2)
B = tf.transpose(B, perm=[0, 1, 3, 2])

# test to verify result: 
for m in range (M):
    for i in range (N):
        for k in range (K):
            for j in range (N):
                assert B[m, i, k, j] == A[m, i, k] * A[m, j, k]

this test passes without errors

Upvotes: 1

Related Questions