Rebuild n SV-decomposed matrices without loops

Question

I have the problem that I need to enforce rank 2 of n 3x3 matrices and I don't want to use loops for it. So I created an SVD of all matrices and set the minimum element of the diagonal matrix to 0.

import numpy as np
from numpy import linalg as la

[u, s, vT] = la.svd(A)
s[:,2] = 0

So when I just had one matrix I would do the following:

sD = np.diag(s)                 # no idea how to solve this for n matrices
Ar2 = u.dot(sD).dot(vT)         # doing the dot products as two dot products 
                                # using np.einsum for n matrices

Ok, so I have the problem to build my diagonal matrices from an (n,3) array. I tried to use np.diag after some reshapes but I don't think that this function can handle the two dimensions of the s-array. A loop would be a possible solution but it is too slow. So, what is the cleanest and quickest way to build my s-matrices into diagonal form respectively calculating the two dot-products with the given information?

# dimensions of arrays:
# A -> (n,3,3)
# u -> (n,3,3)
# s -> (n,3)
# vT -> (n,3,3)

Answer 1

As you suspected, you can use einsum:

Ar2 = np.einsum('ijk,ik,ikl->ijl', u, s, vT)