Should we always pass `rng` to the model?

[torfjelde]

Back in the day, the evaluator for a @model would look like

function demo(rng, model, varinfo, context, sampler)
    ...
end

or something like this.
But when we started making use of contexts more drastically in #249 , this became instead the simpler

function demo(model, varinfo, context)
    ...
end

and instead we provide the rng and sampler argument sometimes through the Samplingcontext(rng, sampler, context), thus providing a clear separation between when we're sampling and when we're evaluating a model.

However, this consequently doesn't allow someone to define models with inherit randomness in them while still preserving determinisim conditional on a rng.

A simple example is an implementation of a model with subsampling of the data. I could do this as follows:

julia> using DynamicPPL, Distributions

julia> @model function demo(y, batch_size=1)
           x ~ Normal()
           y_indices = rand(1:length(y), batch_size)
           y[y_indices] .~ Normal(x, 1)
       end
demo (generic function with 4 methods)

julia> model = demo(randn(32), 4);

julia> model()
4-element view(::Vector{Float64}, [11, 18, 20, 5]) with eltype Float64:
 -0.35070761628093716
  0.19218100395285695
  0.8506289607980133
  1.0317998072662038

julia> rand(model)
(x = -0.991550644289528,)

However, the issue in the above is ofc that the rand call inside @model doesn't have access to the rng used internally in the model, and thus cannot ensure that everything is deterministic given an rng.

Ofc, the user could provide the rng as a specific argument, but this seems quite redundant as we often will have an rng available.

As a result, I'm thinking that it might be useful to remove the rng from the SamplingConctext and instead make it a similarly "private" varriable so users could do

julia> @model function demo(y, batch_size=1)
           x ~ Normal()
           y_indices = rand(__rng__, 1:length(y), batch_size)
           y[y_indices] .~ Normal(x, 1)
       end
demo (generic function with 4 methods)

as they could before.

Thoughts? @mhauru @penelopeysm @yebai @willtebbutt @sunxd3

[mhauru]

I would prefer to force the user to pass the rng themselves as an argument. Reason being that I would like to view all the __x__ variables as implementation details and not part of the user-facing interface. If we do want to expose them (would still prefer to avoid that completely), then I would favour having very simple macros like @rng that would return __rng__, so that we could change the implementation later without breaking the interface.

[willtebbutt]

Ofc, the user could provide the rng as a specific argument, but this seems quite redundant as we often will have an rng available.

Where is it that it will tend to be available? I mean, is it always going to be obvious to a user which rng winds up getting used?

[yebai]

I would favour having very simple macros like @rng

Second, having a Turing.@rng API seems like the most robust approach here.

[torfjelde]

I would prefer to force the user to pass the rng themselves as an argument.

Problem here is that the user then has to specify two rngs when calling a bunch of functions, e.g. sample. Moreover, rng is stateful and will be mutated upon every call to rand. This goes against the notion that a Model shouldn't itself be mutated upon each evaluation.

I would favour having very simple macros like @rng

Regarding the @rng macro, I'm assuming you mean not passing it as an argument, but instead extracting it from the context or returning Random.default_rng(), right?
This is an alternative I did consider, but it won't cut it.

For example, samplers which only use log-density evaluation, e.g. Turing.externalsampler, effectively does the following internally

f = LogDensityFunction(model)
sample(f, external_sampler, ...)

We would then like a user to be able to do the following:

sample(rng, model, external_sampler, ...)

and know that the provided rng is also used within their model. This would not be the case if you used something like @rng. However, if LogDensityFunction also took in the rng as an argument, then this would be easy enough for the external_sampler to address.

Where is it that it will tend to be available? I mean, is it always going to be obvious to a user which rng winds up getting used?

It's used in the tilde pipeline whenever we have a SamplingContext present.

[mhauru]

Problem here is that the user then has to specify two rngs when calling a bunch of functions, e.g. sample.

I don't see this as a problem. The two rngs would serve two quite distinct purposes.

Moreover, rng is stateful and will be mutated upon every call to rand. This goes against the notion that a Model shouldn't itself be mutated upon each evaluation.

I'm not sure this is a problem. Do you have a use case where this would result in unexpected or undesirable behaviour? I'm more tolerant of rngs mutating than mutation more generally. If I do something like

@model function my_model(rng)
    blahblah
end

model_instance = my_model(rng)
model_instance()
sample(model_instance, NUTS(), 100)
rand(model_instance)

then I wouldn't be offended that the results I get depend on things like how many samples I pulled or which sampler I used. Actually, that's exactly what I would expect. I certainly wouldn't want all my samples to be identical.

Regarding the @rng macro, I'm assuming you mean not passing it as an argument, but instead extracting it from the context or returning Random.default_rng(), right?

I was actually thinking of something simpler: If we would come to the conclusion that we want to make __rng__ part of the user-facing interface, then I would just replace it with @rng that returns __rng__. This seems trivial, but

1. maintains a distinction that __x__ names are not for users,
2. is less unexpectedly magical, because one expects some magic with macros, whereas having undefined variable names in scope is quite unexpected.
3. let's us change where the rng lives and how it's passed around without breaking the interface.

[torfjelde]

I don't see this as a problem. The two rngs would serve two quite distinct purposes.

Sure, but there's a reason why we try to stick to using a single RNG in a program.

Simple example: I want to parallelize the sampling within my sampler, so when the user does

sample(rng, model, parallel_sampler, ...)

I convert this into

rngs = split(rng, num_devices)
pmap(rngs) do rng
    sample(rng, model, parallel_sampler, ...)
end

Here we're clearly not achieving what we want if there's also a separate rng contained in model.

I was actually thinking of something simpler: If we would come to the conclusion that we want to make rng part of the user-facing interface, then I would just replace it with @rng that returns rng. This seems trivial, but

Ah, sure then I'm very much whatever. I agree that @rng is better than __rng__, but that raises the question of whether we should do this for all private variables, etc. Personally, don't think it's worth the hasle (the people writing models which explciitly use rng won't be scared of touching __rng__ vs. @rng 🤷 ), but this is not something I have strong opinions about. My quetsion is more about passing rng to model evaluator always or not.