aesara-devs / aehmc

An HMC/NUTS implementation in Aesara

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Re-using the cholesky factor of the mass matrix when computing the kinetic energy of the momentum?

zoj613 opened this issue · comments

aehmc/aehmc/metrics.py

Lines 33 to 34 in 87b0c4d

velocity = matmul(inverse_mass_matrix, momentum)
kinetic_energy = 0.5 * aet.dot(velocity, momentum)

It seems like we could do something like

vecprod = matmul(momentum, tril_inv)
kinetic_energy = 0.5 * dot(vecprod, vecprod)

as an alternative. Would writing it this way cause the graph to be optimized any differently?

Could you elaborate on why you think this would be a better alternative?

I do not particularly think it is better or worse. Remember, this was a question about why/when would one choose one over the other, which you recommended I turn into a github issue. So, I don't have an answer to that question. I was hoping you would let me know.

I was just following the way I write the calculations by hand, never put much thought into that :) I wouldn't expect the performance to be different but maybe you have a different intuition?

Alright, thanks for clarifying. Closing this now.

If it is more efficient to write it one way or another, that might make for a good addition to the graph optimizations.