May I know if it is possible to track the NAN energy during the leapfrog steps in NUTs or HMC?
I don’t recommend this, but in the past I have inserted the following line
print("Energy:", right.energy, " State:" right.q)
right after this code in
def _single_step(self, left, epsilon): """Perform a leapfrog step and handle error cases.""" try: # `State` type right = self.integrator.step(epsilon, left)
This is a really quick way to get a print out of the state and energy. The
q attribute refers to the canonical position coordinate for the model Hamiltonian assumed by the sampler. If you want to print out the momenta variables as well (i.e. to see if too much momentum is sending you into a bad region), you could print out
right.p as well.