TuringLang / TuringLang/DynamicPPL.jl

Should we always pass `rng` to the model?

Open
#721 8 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Julia
Stars
286
Forks
41
Avg merge
1d 11h
Merged PRs (30d)
34

Description

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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing the @model evaluator and SamplingContext paths described in the issue, then reproduce the subsampling example to observe how randomness is currently sourced. A complete change would need an agreed API design that preserves deterministic model behavior for a supplied rng, along with the corresponding implementation and tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.