Is your feature request related to a problem? Please describe.
I used this great package for the first time in a few years, and I realized that it was not clear to me in the documentation whether the callback function f(x; theta) should be vectorized w.r.t x or not. The function should be vectorized looking at the source code and my previous usage. I cannot find a description of the expected output of function f(x; theta) in the documentation.
Describe the solution you'd like
I suggest the documentation states that the expected output of f(x; theta) is a vector of size equal to that of vector input x, and could be written using clearer notation that suggests a vector-valued function with vector-valued arguments. A user does not then need to encounter an error or look through source code to figure out the expected behavior of the callback function.