Improve stacktraces by using custom tag for ForwardDiff #1841

devmotion · 2022-06-20T12:24:29Z

This should improve stacktraces and still be safe. Was suggested by @ChrisRackauckas on discourse.

Ref: https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/

I'll add an example with failing AD in the PR to show if and how it improves stacktraces.

devmotion · 2022-06-20T13:00:04Z

With this PR the example on discourse yields the following stacktrace when failing: https://gist.github.com/devmotion/ab6f8952c576c6ae3a8b65d398939bba#file-stacktrace-txt= Whereas on the master branch one gets: https://pastebin.com/92pGY6Cr

Compare e.g. the information for dot_observe:

On the master branch

[7] dot_observe(dists::Vector{Normal{ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:β, :K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(fit_cucker_smaile), (:data, :cucker_smaile_problem, :problem_p, :global_p), (), (), Tuple{Vector{Float64}, ODEProblem{Matrix{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64}, ODEFunction{true, typeof(cuckersmale!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Int64, Float64, Float64}, Tuple{RK4, Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 3}}}, value::Vector{Float64}, vi::DynamicPPL.ThreadSafeVarInfo{DynamicPPL.TypedVarInfo{NamedTuple{(:β, :K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, 
Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:β, :K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(fit_cucker_smaile), (:data, :cucker_smaile_problem, :problem_p, :global_p), (), (), Tuple{Vector{Float64}, ODEProblem{Matrix{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64}, ODEFunction{true, typeof(cuckersmale!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Int64, Float64, Float64}, Tuple{RK4, Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 3}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:β, 
:K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, 
Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, 
DynamicPPL.Model{typeof(fit_cucker_smaile), (:data, :cucker_smaile_problem, :problem_p, :global_p), (), (), Tuple{Vector{Float64}, ODEProblem{Matrix{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64}, ODEFunction{true, typeof(cuckersmale!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Int64, Float64, Float64}, Tuple{RK4, Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 3}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:β, :K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(fit_cucker_smaile), (:data, :cucker_smaile_problem, :problem_p, :global_p), (), (), Tuple{Vector{Float64}, ODEProblem{Matrix{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64}, ODEFunction{true, typeof(cuckersmale!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Int64, Float64, Float64}, Tuple{RK4, Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 3}}, Vector{Set{DynamicPPL.Selector}}}}}, ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:β, :K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(fit_cucker_smaile), (:data, :cucker_smaile_problem, :problem_p, :global_p), (), (), Tuple{Vector{Float64}, ODEProblem{Matrix{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64}, ODEFunction{true, typeof(cuckersmale!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Int64, Float64, Float64}, Tuple{RK4, Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 3}}, Vector{Base.RefValue{ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:β, :K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(fit_cucker_smaile), (:data, :cucker_smaile_problem, :problem_p, :global_p), (), (), Tuple{Vector{Float64}, ODEProblem{Matrix{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64}, ODEFunction{true, typeof(cuckersmale!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Int64, Float64, Float64}, Tuple{RK4, Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 3}}}})
    @ DynamicPPL C:\Users\jmmsm\.julia\packages\DynamicPPL\e6mZw\src\context_implementations.jl:622

With this PR

  [7] dot_observe(dists::Vector{Normal{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 3}}}, value::Vector{Float64}, vi::DynamicPPL.ThreadSafeVarInfo{DynamicPPL.TypedVarInfo{NamedTuple{(:β, :K, :var), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:β, Setfield.IdentityLens}, Int64}, Vector{Uniform{Float64}}, Vector{AbstractPPL.VarName{:β, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 3}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:K, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:K, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 3}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:var, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:var, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 3}}, Vector{Set{DynamicPPL.Selector}}}}}, ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 3}}, Vector{Base.RefValue{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 3}}}})
    @ DynamicPPL ~/.julia/dev/DynamicPPL/src/context_implementations.jl:622

src/essential/ad.jl

coveralls · 2022-06-20T18:11:00Z

Pull Request Test Coverage Report for Build 2533951369

37 of 37 (100.0%) changed or added relevant lines in 3 files are covered.
3 unchanged lines in 2 files lost coverage.
Overall coverage decreased (-0.3%) to 81.409%

Files with Coverage Reduction	New Missed Lines	%
src/inference/Inference.jl	1	83.87%
src/essential/ad.jl	2	82.19%

Totals
Change from base Build 2523196356:	-0.3%
Covered Lines:	1213
Relevant Lines:	1490

💛 - Coveralls

codecov · 2022-06-21T01:17:10Z

Codecov Report

Merging #1841 (e5a98a8) into master (8adfa22) will decrease coverage by 0.32%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #1841      +/-   ##
==========================================
- Coverage   81.73%   81.40%   -0.33%     
==========================================
  Files          24       24              
  Lines        1500     1490      -10     
==========================================
- Hits         1226     1213      -13     
- Misses        274      277       +3

Impacted Files	Coverage Δ
src/essential/Essential.jl	`100.00% <ø> (ø)`
src/Turing.jl	`100.00% <100.00%> (ø)`
src/essential/ad.jl	`82.19% <100.00%> (-3.89%)`	⬇️
src/essential/compat/reversediff.jl	`94.11% <100.00%> (-1.24%)`	⬇️
src/inference/Inference.jl	`83.87% <0.00%> (-0.81%)`	⬇️

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 8adfa22...e5a98a8. Read the comment docs.

yebai

Thanks, @devmotion - it looks good. Only one minor comment, see below.

yebai · 2022-06-21T17:16:50Z

src/Turing.jl

 end

+# Log density function
+struct LogDensityFunction{V,M,S,C}


This looks great - I think we might want to move it and the callable object function below to DynamicPPL in the future.

That were my thoughts as well. However, I think it would be good to (re-)use it in the optimization module and for DynamicHMC. I assumed it would be easier to do that before moving it to DynamicPPL since possibly it requires some minor changes. I thought it would be cleaner to do that in a follow-up PR and hence did not work on it in this PR.

Sure, sounds good to do it in a follow-up PR.

Improve stacktraces by using custom tag for ForwardDiff

39e88df

devmotion force-pushed the dw/standardtag branch from 02f66c2 to 39e88df Compare June 20, 2022 12:25

devmotion added 2 commits June 20, 2022 14:31

Fix typo

82a26be

Additional fixes

54b74df

devmotion commented Jun 20, 2022

View reviewed changes

src/essential/ad.jl Outdated Show resolved Hide resolved

devmotion commented Jun 20, 2022

View reviewed changes

src/essential/ad.jl Outdated Show resolved Hide resolved

devmotion mentioned this pull request Jun 20, 2022

Concrete terse structs #1830

Closed

Simplify code and define LogDensityFunction

c3624fb

devmotion added 2 commits June 21, 2022 09:36

A bit simpler

96e957f

Add tests

e5a98a8

devmotion requested a review from yebai June 21, 2022 09:40

yebai approved these changes Jun 21, 2022

View reviewed changes

devmotion merged commit 9951638 into master Jun 21, 2022

delete-merged-branch bot deleted the dw/standardtag branch June 21, 2022 18:03

devmotion mentioned this pull request Jun 25, 2022

Unify log density function types #1846

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Improve stacktraces by using custom tag for ForwardDiff #1841

Improve stacktraces by using custom tag for ForwardDiff #1841

Uh oh!

devmotion commented Jun 20, 2022

Uh oh!

devmotion commented Jun 20, 2022 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

coveralls commented Jun 20, 2022 •

edited

Loading

Uh oh!

codecov bot commented Jun 21, 2022 •

edited

Loading

Uh oh!

yebai left a comment

Uh oh!

yebai Jun 21, 2022

Uh oh!

devmotion Jun 21, 2022

Uh oh!

yebai Jun 21, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Improve stacktraces by using custom tag for ForwardDiff #1841

Improve stacktraces by using custom tag for ForwardDiff #1841

Uh oh!

Conversation

devmotion commented Jun 20, 2022

Uh oh!

devmotion commented Jun 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

On the master branch

With this PR

Uh oh!

Uh oh!

Uh oh!

coveralls commented Jun 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 2533951369

💛 - Coveralls

Uh oh!

codecov bot commented Jun 21, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

yebai left a comment

Choose a reason for hiding this comment

Uh oh!

yebai Jun 21, 2022

Choose a reason for hiding this comment

Uh oh!

devmotion Jun 21, 2022

Choose a reason for hiding this comment

Uh oh!

yebai Jun 21, 2022

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

devmotion commented Jun 20, 2022 •

edited

Loading

coveralls commented Jun 20, 2022 •

edited

Loading

codecov bot commented Jun 21, 2022 •

edited

Loading