Einstein Notation for product of stacked matrices

2.1k Views Asked by At

Background Information:

I recently started using the Einstein summation notation to express certain operations over an "image" $\mathbf{A}$ where to each pixel a square matrix is attached. That is, an array of the shape $\left[n_{row}\times n_{col}\times n \times n\right]$. Now, if I want to represent the multiplication of each "submatrix" $A_{pq}$ of size $n \times n$ by a matrix $\mathbf{B}$ of compatible size, I can express it as:

$A^{'}_{pqij} = B_{ik}A_{pqkj}$.

The problem: I have an array $\mathbf{B}$ of the same size as $\mathbf{A}$ and I want to compute a third array $\mathbf{A}^{'}$ where each pixel $\mathbf{A}_{pix}^{'}$ at the first two indices $p$ and $q$ corresponds to the matrix multiplication of matrix $\mathbf{A}_{pix}$ with $\mathbf{B}_{pix}$ in array $\mathbf{B}$ at the same indices $p$ and $q$.

In the notation of linear algebra, for each $n \times n$ matrix indicized by $p$ and $q$ I want:

$A_{pix}^{'} = B_{pix}A_{pix}$

Attempt:

I initially tried something of the form: $A_{mnloij}^{'} = B_{mnik}A_{lokj}$ but if I think this will compute the submatrix multiplication at all combinations of indices $m,l,n,o$.

The Question in Brief: Is there a way, using the Einstein summation notation, to compute the sum:

$A_{pqij}^{'} = B_{mnik}A_{lokj}$

for the combinations of indices where $(m,n)=(l,o)=\left(p,q\right)$ only?

Additional Attempt: In the worst case, I suppose that I have to first compute the expression above and then only slice the array in order to select the "diagonals" where $(m,n)=(l,o)$.

1

There are 1 best solutions below

1
On

Define the tensor

$${E_{ijk}} = \mathrm{1\ \ if\ i=j=k,\ 0\ otherwise}$$

$\mathbf{E}$ has some interesting properties: $$ \mathbf{E.a = Diag(a)} = \mathrm{matrix\ (\mathbf{a}\ along\ the\ diagonal)} $$ $$ \mathbf{E:A = diag(A)} = \mathrm{vector\ (diagonal\ of\ \mathbf{A})}$$ $$ \mathbf{a.E.b = E:ab} = \mathrm{Schur\ product\ (\mathbf{a \odot b})}$$

Your question is similar to the Schur product case

$${A'_{pqij} = E_{pml}\ E_{qno}\ B_{mnik}\ A_{lokj}}$$