title: How to map the kronecker product along array dimensions?
body: <p>Given two tensors A and B with the same dimension <code>(d&gt;=2)</code> and shapes <code>[A_{1},...,A_{d-2},A_{d-1},A_{d}]</code> and <code>[A_{1},...,A_{d-2},B_{d-1},B_{d}]</code> (shapes of the first d-2 dimensions are identical).</p>
  <p>Is there a way to calculate the kronecker product over the last two dimensions?
  The shape of <code>my_kron(A,B)</code>should be <code>[A_{1},...,A_{d-2},A_{d-1}*B_{d-1},A_{d}*B_{d}]</code>.
  For example with <code>d=3</code>,</p>
  <pre><code>A.shape=[2,3,3]
  B.shape=[2,4,4]
  C=my_kron(A,B)
  </code></pre>
  <p><code>C[0,...]</code> should be the kronecker product of <code>A[0,...]</code> and <code>B[0,...]</code> and <code>C[1,...]</code> the kronecker product of <code>A[1,...]</code> and <code>B[1,...]</code>.</p>
  <p>For d=2 this is simply what the <code>jnp.kron</code>(or <code>np.kron</code>) function does.</p>
  <p>For d=3 this can be achived with <code>jax.vmap</code>.
  <code>jax.vmap(lambda x, y: jnp.kron(x[0, :], y[0, :]))(A, B)</code></p>
  <p>But I was not able to find a solution for general (unknown) dimensions.
  Any suggestions?</p>