search

probcomp / Gen.jl

A general-purpose probabilistic programming system with programmable inference
https://gen.dev
Apache License 2.0
1.79k stars 160 forks source link

Add support for user-supplied RNG state in all interfaces #520

Open bgroenks96 opened 7 months ago

bgroenks96 commented 7 months ago

Updates the Gen function interface to use the standard Julia pattern for user-supplied RNG states, i.e:

myfunc(args...) = myfunc(default_rng(), args...)
myfunc(rng::AbstractRNG, args...) = ...

This is applied to all function interfaces which use rng.

Inference algorithms provide instead a keyword argument rng which tends to be more common for higher level function interfaces.

Note that this PR should be fully backwards compatible with all tests and existing Gen code since method dispatches with default_rng() are universally provided.

Resolves #33

fsaad commented 7 months ago

Hi @bgroenks96, thank you for this thorough PR. I support merging these improvements to the API.

Could you please take a look at the failing ContinuousIntegration tests? There are appears to be a "Not implemented" error in one of the tests. https://github.com/probcomp/Gen.jl/actions/runs/8067611177/job/22038710386?pr=520

@alex-lew @ztangent Please also take a look.

bgroenks96 commented 7 months ago

@fsaad Fixed in the last two commits. All tests are passing for me locally.

bgroenks96 commented 7 months ago

Regarding the fix in e39459b, I am not sure that I fully understand this part of the API, but I am assuming by the name that deterministic functions should not need access to the RNG. If this is wrong, then we would unfortunately need to break this part of the API, I think.

ztangent commented 7 months ago

Thanks for this PR! I'll review this in the next couple of days, but intuitively this seems like the right way to modify the interfaces and I've checked that other implementations of Gen (e.g. the work-in-progress JAX implementation) also adopt a similar interface for control over the RNG / RNG seed.

bgroenks96 commented 7 months ago

Thanks for the detailed review, @ztangent . I was afraid this would turn out to be more complicated than it initially looked...

I will try to go through the comments later this week. In the meantime, it would be good to go ahead and create that branch, and then I can change the PR to target this instead of master.

bgroenks96 commented 5 months ago

Hi @ztangent, I apologize for the long delay. I had some other more pressing deadlines to attend to.

I think that I have addressed your first set of comments, as well as the issues with the combinators. I still need to look more closely at the inference algorithms.

Please let me know if I have missed anything or if I did not fully address any of the issues.

EDIT: Note that I have verified that all tests are passing on my machine (as of 0539d93).

ztangent commented 5 months ago

Awesome, thank you! I should have time to look more at this the week after next.

I've also created the modular-rng branch as discussed before, so we can make the changes to the inference library as a separate PR.