Kronecker product like operation between 3D and 2D matrix I'm looking for an operation that does the following computation between a three-dimensional matrix $A$ of shape $T \times T \times J$ and a two-dimensional matrix $B$ shaped $J \times J$ :
$
\DeclareMathOperator{\diag}{diag}
\left[\begin{array}{cccc}
\diag({a}_{1,1}) \mathbf{B} & 
\diag({a}_{1,2})  \mathbf{B} & \cdots & 
\diag({a}_{1, T}) \mathbf{B} \\ 
\diag({a}_{2,1})  \mathbf{B} & 
\diag({a}_{2,2})  \mathbf{B} & \cdots & 
\diag({a}_{2, T}) \mathbf{B} \\ \vdots & \vdots & \vdots & \vdots \\ 
\diag({a}_{T, 1}) \mathbf{B} & 
\diag({a}_{T, 2}) \mathbf{B} & \cdots & 
\diag({a}_{T, T}) \mathbf{B}
\end{array}\right]$
With the resulting matrix having shape $TJ \times TJ$.
Is there a common definition of such an operation, and how would one implement it in Numpy?
In my opinion, this has some similarities to a Kronecker product, which is given
between a $n \times p$ matrix $\mathbf{A}$ and a $m\times p$ matrix $\mathbf{B}$ as follows:
$\mathbf{A} \otimes \mathbf{B}=\left[\begin{array}{cccc}a_{1,1} \mathbf{B} & a_{1,2} \mathbf{B} & \cdots & a_{1, p} \mathbf{B} \\ a_{2,1} \mathbf{B} & a_{2,2} \mathbf{B} & \cdots & a_{2, p} \mathbf{B} \\ \vdots & \vdots & \vdots & \vdots \\ a_{n, 1} \mathbf{B} & a_{n, 2} \mathbf{B} & \cdots & a_{n, p} \mathbf{B}\end{array}\right]$
 A: You could probably implement something like this:
$$(B_{i\%J,j\%J} \cdot A_{i\%J,j\%J,i\%J})_{i,j}$$
where $i \% J$ is the remainder modulo the size $J$.
A: My Python/NumPy is kind of rusty, but in Julia/Matlab I'd do it like this
# Create a random problem
J,T = 3,5
A,B,C = rand(T,T,J), rand(J,J), zeros(J*T,J*T)

# Load the product matrix in B-sized blocks, i.e. JxJ
for i=1:T
  for k=1:T
    C[(1-J+i*J):(i*J), (1-J+k*J):(k*J)] = Diagonal(A[i,k,:])*B
  end
end

A: This probably isn't the most efficient method, but here's one way to go.
Step 1: Build the matrix
$$
M = \DeclareMathOperator{\diag}{diag}
\left[\begin{array}{cccc}
\diag({a}_{1,1})  & 
\diag({a}_{1,2})   & \cdots & 
\diag({a}_{1, T})  \\ 
\diag({a}_{2,1})   & 
\diag({a}_{2,2})   & \cdots & 
\diag({a}_{2, T})  \\ \vdots & \vdots & \vdots & \vdots \\ 
\diag({a}_{T, 1})  & 
\diag({a}_{T, 2})  & \cdots & 
\diag({a}_{T, T}) 
\end{array}\right],
$$
Step 2: Compute the product $M (I_J \otimes B)$.
Here's some code to implement this.
import numpy as np
from numpy.random import randint
import sympy as sp

T = 2
J = 3

A = randint(10,size = (T,T,J))
B = randint(10,size = (J,J))

M = np.zeros((T*J,T*J))
for i in range(T):
    for j in range(J):
        M[J*i + j,j::J] = A[i,:,j]

mat = M@np.kron(np.eye(T),B)


# PRINT ALL MATRICES
        
print('Entries of A:')
for j in range(J):
    sp.pprint(sp.Matrix(A[:,:,j]))
print('Entries of B:')
sp.pprint(sp.Matrix(B))
print('Matrix M:')
sp.pprint(sp.Matrix(M))

print('Final matrix:')
sp.pprint(sp.Matrix(mat))

