Compute Inverse Of 2d Arrays Along The Third Axis In A 3d Array Without Loops

I have an array A whose shape is (N, N, K) and I would like to compute another array B with the same shape where B[:, :, i] = np.linalg.inv(A[:, :, i]). As solutions, I see map and

Solution 1:

For an invertible matrix M we have inv(M).T == inv(M.T) (the transpose of the inverse is equal to the inverse of the transpose).

Since np.linalg.inv is broadcastable, your problem can be solved by simply transposing A, calling inv and transposing the result:

B = np.linalg.inv(A.T).T

For example:

>>> N, K = 2, 3
>>> A = np.random.randint(1, 5, (N, N, K))
>>> A
array([[[4, 2, 3],
        [2, 3, 1]],

       [[3, 3, 4],
        [4, 4, 4]]])

>>> B = np.linalg.inv(A.T).T
>>> B
array([[[ 0.4  , -4.   ,  0.5  ],
        [-0.2  ,  3.   , -0.125]],

       [[-0.3  ,  3.   , -0.5  ],
        [ 0.4  , -2.   ,  0.375]]])

You can check the values of B match the inverses of the arrays in A as expected:

>>> all(np.allclose(B[:, :, i], np.linalg.inv(A[:, :, i])) for i in range(K))
True