MersenneTwister frule

JuliaDiff · Jul 1, 2020 · a746a77 · a746a77
1 parent a5762db
commit a746a77
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/src/rulesets/Random/random.jl b/src/rulesets/Random/random.jl
@@ -1,3 +1,5 @@
+frule(Δargs, ::typeof(MersenneTwister), args...) = MersenneTwister(args...), Zero()
+
 function rrule(::typeof(MersenneTwister), args...)
     function MersenneTwister_rrule(ΔΩ)
         return (NO_FIELDS, map(_ -> Zero(), args)...)

diff --git a/test/rulesets/Random/random.jl b/test/rulesets/Random/random.jl
@@ -1,11 +1,19 @@
 @testset "random" begin
     @testset "MersenneTwister" begin
         @testset "no args" begin
+            rng, dΩ = frule((5.0,), MersenneTwister)
+            @test rng isa MersenneTwister
+            @test dΩ isa Zero
+
             rng, pb = rrule(MersenneTwister)
             @test rng isa MersenneTwister
             @test first(pb(10)) isa Zero
         end
         @testset "unary" begin
+            rng, dΩ = frule((5.0, 4.0), MersenneTwister, 123)
+            @test rng isa MersenneTwister
+            @test dΩ isa Zero
+
             rng, pb = rrule(MersenneTwister, 123)
             @test rng isa MersenneTwister
             @test all(map(x -> x isa Zero, pb(10)))