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
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