Getting the location + gradients of divergences (not of iteration starting points)

So I did a quick and veeeeery dirty hack to get the locations where base_hmc::divergent_ becomes true (ignoring the trajectory). And it looks helpful!

his is the same model as above, green triangles are transitions without divergence, blue squares are the diverging iterations as currently reported, red circles are the “ends” of diverging transitions (points where base_hmc::divergent_ became true). Now the sharp boundary where I would suspect the divergences to happen is clearly visible. And it is IMHO easy to get out of the sampler without any performance hit, the code gets executed only for divergent transitions and is basically the same thing that writes the normal transitions, so even if everything diverged, it will just double the I/O…

2 Likes