docs: def _kp2(A, B):
      """Special case Kronecker tensor product of A[i] and B[i] at each
      time interval i for i = 0 .. N-1
      Specialized for the case A and B rank 3 with A.shape[0]==B.shape[0]
      """
      N = A.shape[0]
      if B.shape[0] != N:
          raise(ValueError)
      newshape1 = A.shape[1]*B.shape[1]
      return np.einsum('ijk,ilm->ijlkm', A, B).reshape(N, newshape1, -1)