Add rewrite to optimize `dot(kron(a, b), c)`

[jessegrabowski]

Description

This can be re-written according to the following relationship:

$$ (A \otimes B) C = \text{vec}(B X A^T) $$

Where $\otimes$ is the kronecker product, and the $\text{vec}$ operation ravels a matrix in column-major order. $X$ is a matrix formed by reshaping $C$ (in column-major order) to conform with $B$. This avoids working with the large kronecker product matrix, and instead gets the result in terms of the much smaller components. Code example:

n = 100
a, b = np.random.normal(size=(2, n, n))
c = np.random.normal(size=(n ** 2, ))

def kronAB_C_clever(a, b, c):
    return (b @ c.reshape((n, n)).T @ a.T).T.ravel()

def direct(a, b, c):
    K = np.kron(a, b)
    x2 = K @ c

%timeit kronAB_C_clever(a, b, c)
73.5 μs ± 3.61 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

%timeit direct(a, b, c)
245 ms ± 72.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

This trick is already used in PyMC here, but only in a limited context. PyMC applies this identity to solve_triangular as well, but it can (and should) also be applied to other types of solve.