Passing state + kwargs to callback

[theogf]

Is there any reason for the callback function to not pass the state as well the additional kwargs... passed to sample?

The motivation for state is obvious (I suppose) and for kwargs... it would allow to pass storage/testing variables more easily.
For example :
mycallback(rng, model, sampler, sample; X_test) = my_test(sample, X_test)

Happy to make a PR if this would be an OK thing to add.

EDIT: Note that this would be breaking.

[cpfiffer]

I think it's probably just a minor oversight that it's not doing this already. I'd be fine doing this.