[Python] Vectorisation in numpy.linalg.tensorsolve

I am trying to solve the equation:

Ax = B

for x where A and B are matrices of shape (3, 3, 3, 3) and (3, 3). This can be handled using numpy.linalg.tensorsolve():

numpy.linalg.tensorsolve(A, B, axes=(0, 2))


I am now trying to solve this equation in n different instances which means my matrices, A and B, are now shaped (10, 3, 3, 3, 3) and (10, 3, 3) respectively.

How do I go about using tensorsolve in this case to keep my code vectorised? I don't know what axes to specify or if it is even possible.

Thanks.

Continue reading...