Fix forkr - Handle rng in Trace #23

FredericWantiez · 2021-04-06T14:21:31Z

Still very early, just opening it to track changes/ideas.

devmotion · 2021-04-06T14:44:41Z

Maybe it would be good to merge #20 first?

FredericWantiez · 2021-04-17T16:38:57Z

Agree, we need to build on #20 for this fix.

yebai · 2021-04-26T20:29:21Z

Agree, we need to build on #20 for this fix.

I've merged #20.

src/rng.jl

src/smc.jl

codecov · 2021-05-17T19:23:43Z

Codecov Report

Merging #23 (7f1b12f) into master (350e2c1) will increase coverage by 4.14%.
The diff coverage is 96.77%.

❗ Current head 7f1b12f differs from pull request most recent head 705d1dc. Consider uploading reports for the commit 705d1dc to get more accurate results

@@            Coverage Diff             @@
##           master      #23      +/-   ##
==========================================
+ Coverage   55.70%   59.85%   +4.14%     
==========================================
  Files           5        6       +1     
  Lines         368      411      +43     
==========================================
+ Hits          205      246      +41     
- Misses        163      165       +2

Impacted Files	Coverage Δ
src/container.jl	`92.36% <94.59%> (-0.27%)`	⬇️
src/rng.jl	`100.00% <100.00%> (ø)`
src/smc.jl	`97.22% <100.00%> (ø)`

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 350e2c1...705d1dc. Read the comment docs.

coveralls · 2021-05-17T19:24:07Z

Pull Request Test Coverage Report for Build 1234943898

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

For more information on this, see Tracking coverage changes with pull request builds.
To avoid this issue with future PRs, see these Recommended CI Configurations.
For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

60 of 62 (96.77%) changed or added relevant lines in 3 files are covered.
1 unchanged line in 1 file lost coverage.
Overall coverage increased (+3.9%) to 59.611%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/container.jl	35	37	94.59%

Files with Coverage Reduction	New Missed Lines	%
src/container.jl	1	91.67%

Totals
Change from base Build 946932953:	3.9%
Covered Lines:	245
Relevant Lines:	411

💛 - Coveralls

…S.jl into feature/traced_rng

FredericWantiez · 2021-06-07T17:28:44Z

@yebai I think there's an issue with using fork and passing the RNG to Trace. When we fork the particles during the resampling step, the RNG is set to the RNG we used in the original Trace definition, all children particles will use it and we won't be able to reconstruct the stream.

Also, on the side, when do we populate the reference particle when we resample in PG ?

yebai

We'll need a mechanism to split the random stream during the resampling step. The splitting operation

allows the children particles to share the random history before the splitting point,
make the random streams produce independent random numbers after the splitting point.

More concretely, we need to split the random stream when copying a trace, i.e. fork(trace):

AdvancedPS.jl/src/container.jl

Line 33 in 156afff

Base.copy(trace::Trace) = Trace(trace.f, copy(trace.ctask))

We need to replay the random stream when replaying a trace, i.e. forkr(trace):

AdvancedPS.jl/src/container.jl

Line 58 in 156afff

function forkr(trace::Trace)

Also, on the side, when do we populate the reference particle when we resample in PG ?

I'm not sure I understand the question. The reference particle is incrementally replayed during the forward propagation of a conditional SMC step. If the reference particle has multiple children at a resampling step, one child will inherit the random stream and keep replaying, the other children will split the random stream and start producing new random numbers.

Also, happy to discuss more when we meet.

yebai · 2021-06-08T16:52:40Z

src/rng.jl

+
+
+# Set seed manually, for init ?
+Random.seed!(rng::TracedRNG, seed) = Random.seed!(rng.rng, seed)


Do we store the seed somewhere in TracedRNG when calling this function?

devmotion · 2021-06-08T20:37:14Z

We'll need a mechanism to split the random stream during the resampling step. The splitting operation
1. allows the children particles to share the random history before the splitting point,

2. make the random streams produce independent random numbers after the splitting point.

Isn't this exactly what splittable RNGs are designed for (e.g. https://github.com/google/jax/blob/master/design_notes/prng.md#design, https://dl.acm.org/doi/abs/10.1145/2660193.2660195, and https://publications.lib.chalmers.se/records/fulltext/183348/local_183348.pdf)? The last reference also discusses why naive approaches for splitting linear PRNGs fail or make it hard to assess the quality of the returned random numbers (I am not sufficiently familiar with this area to determine if these all are valid concerns).

yebai · 2021-06-08T21:08:36Z

yes, we are planning to use splittable RNGs.

…S.jl into feature/traced_rng

FredericWantiez · 2021-07-29T20:20:11Z

Looked at the random module in tensorflow recently, they also use counter-based rng and the split api: https://github.com/tensorflow/tensorflow/blob/v2.5.0/tensorflow/python/ops/stateful_random_ops.py#L933-L981

We could do something similar:

mutable struct TracedRNG{R,T<:Random123.AbstractR123{R}}
    count::Int
    rng::T
    keys::Array{R}
end


TracedRNG(r::Random123.AbstractR123) = TracedRNG(0, r, [r.key])


function split(r::TracedRNG, n::Integer)
    rngs = map(1:n) do i
        new_seed = hash(r.rng.key, r.rng.ctr1+i)
        new_rng = TracedRNG(r.count, Philox2x(new_seed), [r.keys..., new_seed])
        return new_rng
    end
end

rng = TracedRNG(Philox2x())
a, b = split(rng, 2)

To replay the random numbers we should be able to cycle through the keys and reset the counters

Reset rng

Feature/split

FredericWantiez · 2021-08-29T13:26:06Z

@yebai I simplified the API slightly, we don't need to track both the model counter and the rng counters. We can use the model counter if every particle has an unique key/id. Fixed the naming and the resampling step that was wrong. Also added a few tests. I can add some more, don't think we catch all the edge cases.

yebai

Thanks @FredericWantiez - excellent work. I have gone through the rng.jl file and made some comments along the process. I'll take a look at other parts slightly later.

In particular, I find the function pair split and update_rng! is slightly confusing while reading the code. Given that the split and update_rng! functions are always used in pairs, I wonder whether we can consider merging these functions into split!(rng::TracedRNG)::TracedRNG for better clarity. This split! function should always split a given RNG once, mutate it in place and return the mutated RNG. If we need to split N-times, we simply call split N-times repetitively. Below is an example implementation of the split! function:

function split!(rng::TracedRNG{T}) where {T}
    key = hash(r.rng.key, convert(UInt, r.rng.ctr1)
    seed!(rng, key)
    return set_counter!(rng.rng, rng.count + 1) # rng.count + 1 is used instead of rng.count
end

In order to ensure that calling split! repetitively with the same RNG will return new RNGs with different keys, we increase the internal counter by one every time a split is performed. This way, the next time when split! is called on the same RNG, it will get a different key because r.rng.ctr1 will be different.

A side question: do we actually need to set rng.rng.ctr1 to rng.count when performing a split? If so, is there a reference on this?

src/container.jl

src/rng.jl

Fix naming

FredericWantiez · 2021-09-09T22:51:41Z

@devmotion @yebai Fixed the last few thing we discussed last week, let me know what you think.

src/container.jl

src/rng.jl

devmotion · 2021-09-09T23:07:46Z

src/rng.jl

+import Random123: set_counter!
+
+# Default RNG type for when nothing is specified
+_BASE_RNG = Philox2x


Suggested change

_BASE_RNG = Philox2x

const _BASE_RNG = Philox2x

devmotion · 2021-09-09T23:14:38Z

src/rng.jl

+
+Set key and counter of inner RNG to key and the running model step
+"""
+function seed!(rng::TracedRNG{T}, key) where {T}


Suggested change

function seed!(rng::TracedRNG{T}, key) where {T}

function seed!(rng::TracedRNG, key)

devmotion · 2021-09-09T23:14:58Z

src/rng.jl

+
+Track current key of the inner RNG
+"""
+function save_state!(r::TracedRNG{T}) where {T}


Suggested change

function save_state!(r::TracedRNG{T}) where {T}

function save_state!(r::TracedRNG)

devmotion · 2021-09-09T23:16:10Z

src/rng.jl

+    return push!(r.keys, r.rng.key)
+end
+
+Base.copy(r::TracedRNG{T}) where {T} = TracedRNG(r.count, copy(r.rng), copy(r.keys))


Suggested change

Base.copy(r::TracedRNG{T}) where {T} = TracedRNG(r.count, copy(r.rng), copy(r.keys))

Base.copy(r::TracedRNG) = TracedRNG(r.count, copy(r.rng), copy(r.keys))

I am also worried that copy(r.rng) is not sufficient to ensure that the copy is decoupled from the original RNG. Maybe it would be safer (albeit more expensive) to use deepcopy?

src/rng.jl

test/smc.jl

cpfiffer

I think David's already added most of the things I would add -- in general this seems to be fine with a few minor changes needed.

src/rng.jl

Fix types, import and PR review

src/rng.jl

yebai · 2021-09-14T19:18:30Z

Excellent work - many thanks @FredericWantiez @devmotion and @cpfiffer!

Perhaps we can wait for a few days to see whether we miss anything, then make a new release.

devmotion · 2021-09-14T20:36:27Z

src/container.jl

@@ -1,31 +1,43 @@
-struct Trace{F}
+struct Trace{F,U,N,V<:Random123.AbstractR123{U}}


Do we dispatch on U, N or V somewhere? Otherwise (and probably even in this case) this could be simplified to

struct Trace{F,R<:TracedRNG} f::F ctask::Libtask.CTask rng::R end

devmotion · 2021-09-14T20:37:38Z

src/container.jl

 - consume(pc::ParticleContainer): return incremental likelihood
 """
-mutable struct ParticleContainer{T<:Particle}
+mutable struct ParticleContainer{T<:Particle,U,N,V<:Random123.AbstractR123{U}}


Same here, I guess we could use

mutable struct ParticleContainer{T<:Particle,R<:TracedRNG} vals::Vector{T} logWs::Vector{Float64} rng::R end

?

devmotion · 2021-09-14T20:38:30Z

src/container.jl

+        k = split(pi.rng.rng.key)
+        Random.seed!(pi.rng, k[1])
+    end
+end


I guess we could add an empty return or nothing at the end to avoid that the last seed or rng is returned.

devmotion · 2021-09-14T20:39:48Z

test/container.jl

        function fpc(logp)
            f = let logp = logp
-                () -> begin
+                (rng) -> begin


The brackets could be omitted.

yebai mentioned this pull request Apr 26, 2021

SMC and PG algorithms #20

Merged

FredericWantiez force-pushed the feature/traced_rng branch from 5366257 to da48dab Compare April 30, 2021 10:34

Init

8725b79

FredericWantiez force-pushed the feature/traced_rng branch from 8d50100 to 8725b79 Compare April 30, 2021 11:06

FredericWantiez added 3 commits April 30, 2021 12:08

mcmc chains

81603bd

rng

2382874

Trace

79d9857

devmotion reviewed May 4, 2021

View reviewed changes

src/rng.jl Outdated Show resolved Hide resolved

devmotion reviewed May 4, 2021

View reviewed changes

src/smc.jl Outdated Show resolved Hide resolved

FredericWantiez added 2 commits May 12, 2021 18:03

Rng, test

ddbbc3c

Fix tests

87559ef

FredericWantiez added 4 commits May 20, 2021 15:09

Global rng

460f26c

Format, duplicate rng

b5240eb

Tracedrng

f125369

Merge branch 'master' of https://github.com/FredericWantiez/AdvancedP…

156afff

…S.jl into feature/traced_rng

yebai requested changes Jun 8, 2021

View reviewed changes

FredericWantiez added 4 commits June 13, 2021 18:15

Format

cd12c8c

Track Particle container

5cbbe40

Merge branch 'master' of https://github.com/FredericWantiez/AdvancedP…

2fb7290

…S.jl into feature/traced_rng

Test, format

22fa456

Reset RNG, save state, random123

8f05881

FredericWantiez added 7 commits August 17, 2021 19:01

Downcast

5b01908

Merge pull request #2 from FredericWantiez/feature/split

8134751

Reset rng

Correct replaying mechanism

0893eac

Format

0ecdcce

Fix resampling step

f6df82d

Clean

e76747e

Merge pull request #3 from FredericWantiez/feature/split

d77876f

Feature/split

yebai changed the title ~~[WIP] Fix forkr - Handle rng in Trace~~ Fix forkr - Handle rng in Trace Aug 30, 2021

yebai requested changes Aug 30, 2021

View reviewed changes

src/container.jl Show resolved Hide resolved

src/container.jl Outdated Show resolved Hide resolved

src/rng.jl Outdated Show resolved Hide resolved

src/rng.jl Outdated Show resolved Hide resolved

src/rng.jl Outdated Show resolved Hide resolved

FredericWantiez added 5 commits September 4, 2021 14:58

Fix naming

4cd74b0

Format

3acd49a

Test

fc31d02

Doc string

d709319

Merge pull request #4 from FredericWantiez/feature/split

6979bbe

Fix naming

devmotion reviewed Sep 9, 2021

View reviewed changes

yebai requested a review from cpfiffer September 10, 2021 09:09

cpfiffer reviewed Sep 10, 2021

View reviewed changes

src/rng.jl Outdated Show resolved Hide resolved

FredericWantiez added 2 commits September 11, 2021 13:12

Fix types, import and PR review

8ed2100

Merge pull request #5 from FredericWantiez/feature/split

4f8ec74

Fix types, import and PR review

yebai reviewed Sep 14, 2021

View reviewed changes

src/rng.jl Outdated Show resolved Hide resolved

yebai added 2 commits September 14, 2021 20:05

Update src/rng.jl

7f1b12f

Update Project.toml

705d1dc

yebai merged commit 51c18d4 into TuringLang:master Sep 14, 2021

yebai mentioned this pull request Sep 14, 2021

forkr is likely broken for particle Gibbs #21

Closed

devmotion reviewed Sep 14, 2021

View reviewed changes

FredericWantiez mentioned this pull request Sep 25, 2021

TracedRNG type #31

Merged



		# Set seed manually, for init ?
		Random.seed!(rng::TracedRNG, seed) = Random.seed!(rng.rng, seed)

	function seed!(rng::TracedRNG{T}, key) where {T}
	function seed!(rng::TracedRNG, key)

	function save_state!(r::TracedRNG{T}) where {T}
	function save_state!(r::TracedRNG)

	Base.copy(r::TracedRNG{T}) where {T} = TracedRNG(r.count, copy(r.rng), copy(r.keys))
	Base.copy(r::TracedRNG) = TracedRNG(r.count, copy(r.rng), copy(r.keys))

		@@ -1,31 +1,43 @@
		struct Trace{F}
		struct Trace{F,U,N,V<:Random123.AbstractR123{U}}

Fix forkr - Handle rng in Trace #23

Fix forkr - Handle rng in Trace #23

Uh oh!

Conversation

FredericWantiez commented Apr 6, 2021

Uh oh!

devmotion commented Apr 6, 2021

Uh oh!

FredericWantiez commented Apr 17, 2021

Uh oh!

yebai commented Apr 26, 2021

Uh oh!

Uh oh!

Uh oh!

codecov bot commented May 17, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

coveralls commented May 17, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 1234943898

Warning: This coverage report may be inaccurate.

Details

💛 - Coveralls

Uh oh!

FredericWantiez commented Jun 7, 2021

Uh oh!

yebai left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

devmotion commented Jun 8, 2021

Uh oh!

yebai commented Jun 8, 2021

Uh oh!

FredericWantiez commented Jul 29, 2021

Uh oh!

FredericWantiez commented Aug 29, 2021

Uh oh!

yebai left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

FredericWantiez commented Sep 9, 2021

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

cpfiffer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yebai commented Sep 14, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented May 17, 2021 •

edited

Loading

coveralls commented May 17, 2021 •

edited

Loading

yebai left a comment •

edited

Loading