Is it possible to get the number of potential function evaluation and how many times we perform gradient evaluations during the training? (since in such as NUTS; and all the rejections in warmup; it’s impractical to trace the number of evaluations manually )
You can print/record this number: num_proposals. Right now, this number together with other infos (such as tree_depth, potential_energy, diverging,…) are not recorded during the sampling, so you need to change the source code.
It does not make sense to expose this to the end user since most of them shouldn’t be concerned with it. As @fehiepsi mentioned, you can always print this from the source code. If you want only the numbers in aggregate, you could also just run your code using
cProfile to see how many times is the potential_energy function and grad computation function in the integrator being hit.