sschoenholz

Good question! Since, in the infinite width limit, the output of the network is a gaussian process you can draw samples from the function posterior at points `X = (x_1,...

Great question! We move into a diagonal basis before applying the exponential in which case elementwise exponential and matrix exponential agree. In other words, we write $\exp(At)v = U\exp(Dt)U^Tv$ where...

Hey! Thanks for pushing on this. We'd love to iterate on this to get it working for you (though looking at the UNet architecture I am a bit concerned that...

Hey! I looked into it a bit and I think the approach you're describing should work. However, there might be some subtlety involved in making sure all of the shapes...

Good question! I'm not sure I have a totally satisfactory answer, but my guess is that the function-space -> parameter space map induced by the Jacobian breaks down when the...

Glancing at this code, it seems like there are a lot of per-triangle / per-vertex operations that are likely the origin of the slow behavior (and, to me, it makes...

You can even drop the in_axes, `vmap(partialFunc)(triSimplicesInterior, triNLinterior)`.

+1 to James. @apal90 Perhaps you can set up an example in a colab notebook? I'd be happy to take a glance at it.