I think i got it on my own, but nevertheless will share the solution:
since getting only the diagonals of a matrix multiplication
> Z = N.diag(X.dot(Y))
is equivalent to the individual sum of the scalar product of rows of X and columns of Y, the previous statement is equivalent to:
> Z = (X * Y.T).sum(-1)
For the original variables this means:
> result = (A.dot(B) * A).sum(-1)
Please correct me if I am wrong but this should be it …