Something seems wrong here, you have
\theta = \Lambda(A + XB)
where A is 1 \times 7420 and B is 53 \times7420. The only way this works out is when X is 1 \times 53 in which case \theta is 1 \times 7420; but in the image it is 1 \times 70345? Is your X_idx doing what you think it is? Is there an extra axis being summed over?