DynamicPPL 0.36

Author: penelopeysm

Project.toml

@@ -16,4 +16,4 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-DynamicPPL = "0.35"
+DynamicPPL = "0.36"

README.md

@@ -26,10 +26,6 @@ Note that the links-to-existing-GitHub-issues in the table are also defined in t
 The workflow is the most complicated part of this repository.
 This section attempts to explain it from the 'bottom up'; if you prefer a 'top down' approach start by looking at the GitHub Action workflow, `.github/workflows/test.yml`.
 
-Firstly, there is library code for running the benchmarks.
-This is in `lib.jl`; it should (in the near future) be put directly into DynamicPPL.jl.
-Until then, it has to live here.
-
 Under the hood, the main thing that actually runs the AD tests / benchmarks is `main.jl`.
 You can run `julia --project=. main.jl` and it will print some usage information.
 However, it is the Python script `ad.py` that controls how this Julia script is called.

lib.jl

@@ -1,5 +1,6 @@
 import Test: @test, @testset
 using DynamicPPL: DynamicPPL, VarInfo
+using DynamicPPL.TestUtils.AD: run_ad, ADResult, ADIncorrectException
 using ADTypes
 using Printf: @printf
 
@@ -26,11 +27,6 @@ ADTYPES = Dict(
 include("models.jl")
 using .Models: MODELS
 
-# Benchmarking code is defined here. In time this will be put into DynamicPPL.
-# See https://github.com/TuringLang/DynamicPPL.jl/pull/882
-include("lib.jl")
-using .Lib: run_ad, ADIncorrectException
-
 # The entry point to this script itself begins here
 if ARGS == ["--list-model-keys"]
     foreach(println, sort(collect(keys(MODELS))))
@@ -39,22 +35,37 @@ elseif ARGS == ["--list-adtype-keys"]
 elseif length(ARGS) == 3 && ARGS[1] == "--run"
     model, adtype = MODELS[ARGS[2]], ADTYPES[ARGS[3]]
 
-    if ARGS[2] == "control_flow"
-        # https://github.com/TuringLang/ADTests/issues/4
-        vi = DynamicPPL.unflatten(VarInfo(model), [0.5, -0.5])
-        params = [-0.5, 0.5]
-        result = run_ad(model, adtype; varinfo=vi, params=params, benchmark=true)
-    else
-        result = run_ad(model, adtype; benchmark=true)
-    end
-
-    if isnothing(result.error)
+    try
+        if ARGS[2] == "control_flow"
+            # https://github.com/TuringLang/ADTests/issues/4
+            vi = DynamicPPL.unflatten(VarInfo(model), [0.5, -0.5])
+            params = [-0.5, 0.5]
+            result = run_ad(model, adtype; varinfo=vi, params=params, benchmark=true)
+        else
+            result = run_ad(model, adtype; benchmark=true)
+        end
+        # If reached here - nothing went wrong
         @printf("%.3f", result.time_vs_primal)
-    elseif result.error isa ADIncorrectException
-        println("wrong")
-    else
-        # some other error happened
-        println("error")
+    catch e
+        if result.error isa ADIncorrectException
+            # First check for completely incorrect ones
+            for (a, b) in zip(result.grad_expected, result.grad_actual)
+                if !isnan(a) && !isnan(b) && abs(a - b) > 1e-6
+                    println("wrong")
+                    exit()
+                end
+            end
+            # If not, check for NaN's and report those
+            if any(isnan, result.grad_expected) || any(isnan, result.grad_actual)
+                println("NaN")
+            else
+                # Something else went wrong, shouldn't happen
+                println("wrong")
+            end
+        else
+            # Some other error, just say it's an error
+            println("error")
+        end
     end
 else
     println("Usage: julia main.jl --list-model-keys")