There might be a way to do it with a for loop, and that would definitely be preferable. I gotta think about it some more though to be sure. One thing is that I’ve seen papers (can’t remember what) say that these Kronecker models start to not work so well when the dimensionality is greater than ~6. One possibility is to just code in each case and disallow dimensions > 6. Probably really uncommon. I’m sure two dimensions is very common, maybe three, and then maybe four starts to push it.
Looks like your fix does the trick. Yeah, it’s probably not going to get much faster (unless you can run on the GPU…? which may or may not work). If Eigh is the bottleneck then that’s pretty great. Have you seen MarginalKron? Also this PR for non-gaussian likelihoods here.
Edit: Saw your post on the MarginalKron PR, so yup you saw. Will leave the link for ppl reading.