deepcopy of Modules not supported

[tpapp]

I am trying to add support for AD via Zygote for LogDensityProblems in this PR: tpapp/LogDensityProblems.jl#24, but this line (currently skipped) fails with

julia> @test logdensity(ValueGradient, ∇ℓ, x) ≅ ValueGradient(f(x), -6 .* x)
Error During Test at REPL[28]:1
  Test threw exception
  Expression: logdensity(ValueGradient, ∇ℓ, x) ≅ ValueGradient(f(x), -6 .* x)
  Compiling Tuple{typeof(logdensity),Type{Value},TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity,1},typeof(f)},Array{Float64,1}}: deepcopy of Modules not supported
  Stacktrace:
   [1] error(::String) at ./error.jl:33
   [2] deepcopy_internal(::Module, ::IdDict{Any,Any}) at ./deepcopy.jl:34
   [3] _deepcopy_array_t(::Any, ::Type, ::IdDict{Any,Any}) at ./deepcopy.jl:91
   [4] deepcopy_internal(::Array{Any,1}, ::IdDict{Any,Any}) at ./deepcopy.jl:78
   [5] deepcopy_internal(::Any, ::IdDict{Any,Any}) at ./deepcopy.jl:67
   [6] _deepcopy_array_t(::Any, ::Type, ::IdDict{Any,Any}) at ./deepcopy.jl:91
   [7] deepcopy_internal at ./deepcopy.jl:78 [inlined]
   [8] deepcopy at ./deepcopy.jl:28 [inlined]
   [9] Core.Compiler.IRCode(::IRTools.Meta) at /home/tamas/.julia/packages/IRTools/yAuDJ/src/ir/wrap.jl:104
   [10] _lookup_grad(::Type) at /home/tamas/.julia/packages/Zygote/Ohw1K/src/compiler/emit.jl:124
   [11] #s53#851 at /home/tamas/.julia/packages/Zygote/Ohw1K/src/compiler/interface2.jl:17 [inlined]
   [12] #s53#851(::Any, ::Any, ::Any) at ./none:0
   [13] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any,N} where N) at ./boot.jl:522
   [14] #3 at /home/tamas/code/julia/LogDensityProblems/src/LogDensityProblems.jl:237 [inlined]
   [15] (::Zygote.Pullback{Tuple{getfield(LogDensityProblems, Symbol("##3#4")){TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity,1},typeof(f)}},Array{Float64,1}},Tuple{getfield(Zygote, Symbol("##222#back#144")){getfield(Zygote, Symbol("##142#143")){Zygote.Context,getfield(LogDensityProblems, Symbol("##3#4")){TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity,1},typeof(f)}},Symbol,TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity,1},typeof(f)}}},Zygote.Pullback{Tuple{typeof(logdensity),Type{Value},TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity,1},typeof(f)},Array{Float64,1}},Tuple{typeof(logdensity)}},getfield(Zygote, Symbol("##234#back#147")){getfield(Zygote, Symbol("#back#146")){:value,Zygote.Context,Value{Float64},Float64}}}})(::Int8) at /home/tamas/.julia/packages/Zygote/Ohw1K/src/compiler/interface2.jl:0
   [16] #66 at /home/tamas/.julia/packages/Zygote/Ohw1K/src/compiler/interface.jl:38 [inlined]
   [17] logdensity(::Type{ValueGradient}, ::LogDensityProblems.ZygoteGradientLogDensity{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity,1},typeof(f)}}, ::Array{Float64,1}) at /home/tamas/code/julia/LogDensityProblems/src/AD_Zygote.jl:20
   [18] top-level scope at none:0
   [19] eval(::Module, ::Any) at ./boot.jl:328
   [20] eval_user_input(::Any, ::REPL.REPLBackend) at /home/tamas/src/julia-git/usr/share/julia/stdlib/v1.2/REPL/src/REPL.jl:85
   [21] run_backend(::REPL.REPLBackend) at /home/tamas/.julia/packages/Revise/gStbk/src/Revise.jl:771
   [22] (::getfield(Revise, Symbol("##58#60")){REPL.REPLBackend})() at ./task.jl:259
ERROR: There was an error during testing

I realize the example is far from minimal, but if I am doing something obviously wrong, I would appreciate advice about it.

[MikeInnes]

Definitely a compiler bug. If you can paste an MRP that'd be very helpful, if not a project/manifest with a reproducer script would be great.

[tpapp]

Repo with manifest:
https://github.com/tpapp/ZygoteLogDensityProblemsBug.jl

[tpapp]

Or a standalone script:

using Pkg
pkg"add LogDensityProblems#tp/zygote"
pkg"add Zygote#master IRTools#master"
pkg"add TransformVariables"
using LogDensityProblems, TransformVariables
import Zygote

f(x) = -3*abs2(x[1])
ℓ = TransformedLogDensity(as(Array, asℝ, 1), f)
∇ℓ = ADgradient(:Zygote, ℓ)
x = randn(1)
logdensity(LogDensityProblems.ValueGradient, ∇ℓ, x)

[tpapp]

I updated the repo, the issue still persists.

[tpapp]

Bump, I still have the issue, updated the repo above.

[MikeInnes]

#87 fixes your immediate issue but unfortunately isn't enough to get this example to run. The issue is this function since we don't support try/catch blocks yet (#88). If you refactor that to use return values it'll likely be both faster and more ADable in one go.

[tpapp]

Thanks for the fix and the advice.