Memory allocation limit for NUTS with custom `logp` function (but not with VI methods)

Yes your intuition is mostly correct, except you don’t necessarily need to test the function with random inputs (although sometimes the bottlenecks can be coordinate specific)

You can get a useable point via model.initial_point(). Evaluating the function a couple hundred times on this point should give you a stable profile.

The underscore names are for transformed parameters. Positive distributions will by default have a log transform and the variable name will be f"{name}_log__", other distributions have different default transforms (e.g, simplex for Dirichlet), and some will have none like the Normal, as any value has a nonzero density.

If you want to generate multiple points you can set initval="prior" for each variable and everytime you call model.initial_point you’ll get a new point. But again, a single initial point is usually a good enough for benchmarking.

1 Like