Merge pull request #38 from TuringLang/mt/fix_zygote_segfault

Call Zygote.refresh every 10 distributions in tests

Author: yebai

.github/workflows/ForwardDiff_Tracker.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: [1.0.5, 1.2.0, 1.3]
+        julia-version: [1.0, 1.3]
         julia-arch: [x64, x86]
         os: [ubuntu-latest, macOS-latest]
         exclude:

.github/workflows/Others.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: [1.0.5, 1.2.0, 1.3]
+        julia-version: [1.0, 1.3]
         julia-arch: [x64, x86]
         os: [ubuntu-latest, macOS-latest]
         exclude:

.github/workflows/Zygote.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: [1.0.5, 1.2.0, 1.3]
+        julia-version: [1.0, 1.3]
         julia-arch: [x64, x86]
         os: [ubuntu-latest, macOS-latest]
         exclude:

test/others.jl

@@ -132,7 +132,7 @@ if get_stage() in ("Others", "all")
         d1 = TuringDiagMvNormal(zeros(10), sigmas)
         d2 = MvNormal(zeros(10), sigmas)
 
-        @test entropy(d1) == entropy(d2)
+        @test isapprox(entropy(d1), entropy(d2), rtol = 1e-6)
     end
 
     @testset "Params" begin

test/test_utils.jl

@@ -147,6 +147,8 @@ function get_stage()
     return "all"
 end
 
+const zygote_counter = Ref(0)
+
 function test_ad(f, at = 0.5; rtol = 1e-8, atol = 1e-8)
     stg = get_stage()
     if stg == "all"
@@ -177,7 +179,11 @@ function test_ad(f, at = 0.5; rtol = 1e-8, atol = 1e-8)
             @test isapprox(reverse_tracker, finite_diff, rtol=rtol, atol=atol)
         end
     elseif stg == "Zygote"
+        zygote_counter[] += 1
         isarr = isa(at, AbstractArray)
+        if mod(zygote_counter[], 50) == 0
+            Zygote.refresh()
+        end
         reverse_zygote = Zygote.gradient(f, at)[1]
         if isarr
             forward = ForwardDiff.gradient(f, at)