Support for println / IO in general?

[penelopeysm]

This isn't Turing.jl test-suite related, but I often stick @show inside DynamicPPL functions to see what's going on, and that usually leads to errors if I try to run sampling with AD.

The example below errors, and the same happens with println(x). Would it be possible to support this, or is it too awkward?

using DifferentiationInterface
import Mooncake

function f(x)
    @show x
    return x
end
backend = AutoMooncake(; config=nothing)
prep = prepare_gradient(f, backend, 1.0)

Traceback on Julia 1.10.5, macOS, [email protected], [email protected]:

julia> prep = prepare_gradient(f, backend, 1.0)
x = 1.0
┌ Warning: Unable to put rule in rule field. Rule should error.
└ @ Mooncake ~/.julia/packages/Mooncake/PS8NG/src/interpreter/s2s_reverse_mode_ad.jl:1487
MethodInstance is
MethodInstance for Base.var"#sprint#501"(::Nothing, ::Int64, ::typeof(sprint), ::Function, ::Float64)

with signature
Tuple{Base.var"##sprint#501", Nothing, Int64, typeof(sprint), Function, Float64}

derived_rule is of type
Mooncake.DerivedRule{Tuple{Base.var"##sprint#501", Nothing, Int64, typeof(sprint), Function, Tuple{Float64}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{Mooncake.CoDual{Base.var"##sprint#501", Mooncake.NoFData}, Mooncake.CoDual{Nothing, Mooncake.NoFData}, Mooncake.CoDual{Int64, Mooncake.NoFData}, Mooncake.CoDual{typeof(sprint), Mooncake.NoFData}, Mooncake.CoDual, Mooncake.CoDual{Tuple{Float64}, Mooncake.NoFData}}, Union{}}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{Mooncake.NoRData}, Tuple{Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, Any, Tuple{Float64}}}}, Val{true}, Val{6}}

Expected type is
Mooncake.DerivedRule{Tuple{Base.var"##sprint#501", Nothing, Int64, typeof(sprint), Function, Tuple{Float64}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{Mooncake.CoDual{Base.var"##sprint#501", Mooncake.NoFData}, Mooncake.CoDual{Nothing, Mooncake.NoFData}, Mooncake.CoDual{Int64, Mooncake.NoFData}, Mooncake.CoDual{typeof(sprint), Mooncake.NoFData}, Mooncake.CoDual, Mooncake.CoDual{Tuple{Float64}, Mooncake.NoFData}}, Mooncake.CoDual{String, Mooncake.NoFData}}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{Mooncake.NoRData}, Tuple{Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, Any, Tuple{Float64}}}}, Val{true}, Val{6}}

ERROR: UndefVarError: `args` not defined
Stacktrace:
  [1] _build_rule!(rule::Mooncake.LazyDerivedRule{Tuple{…}, Mooncake.DerivedRule{…}})
    @ Mooncake ~/.julia/packages/Mooncake/PS8NG/src/interpreter/s2s_reverse_mode_ad.jl:1500
  [2] LazyDerivedRule
    @ ~/.julia/packages/Mooncake/PS8NG/src/interpreter/s2s_reverse_mode_ad.jl:1506 [inlined]
  [3] f
    @ ./REPL[41]:2 [inlined]
  [4] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
    @ Base.Experimental ./<missing>:0
  [5] (::MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{…}, Mooncake.CoDual{…}}})(::Mooncake.CoDual{typeof(f), Mooncake.NoFData}, ::Mooncake.CoDual{Float64, Mooncake.NoFData})
    @ MistyClosures ~/.julia/packages/MistyClosures/4rFSA/src/MistyClosures.jl:22
  [6] DerivedRule
    @ ~/.julia/packages/Mooncake/PS8NG/src/interpreter/s2s_reverse_mode_ad.jl:739 [inlined]
  [7] __value_and_pullback!!(::Mooncake.DerivedRule{…}, ::Float64, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
    @ Mooncake ~/.julia/packages/Mooncake/PS8NG/src/interface.jl:14
  [8] value_and_pullback!!(::Mooncake.DerivedRule{Tuple{…}, MistyClosures.MistyClosure{…}, MistyClosures.MistyClosure{…}, Val{…}, Val{…}}, ::Float64, ::Function, ::Float64)
    @ Mooncake ~/.julia/packages/Mooncake/PS8NG/src/interface.jl:117
  [9] value_and_pullback(::Function, ::DifferentiationInterfaceMooncakeExt.MooncakeOneArgPullbackPrep{Float64, Mooncake.DerivedRule{…}, Float64, Float64}, ::AutoMooncake{Nothing}, ::Float64, ::Tuple{Bool})
    @ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/tBGja/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:36
 [10] prepare_pullback(::Function, ::AutoMooncake{Nothing}, ::Float64, ::Tuple{Bool})
    @ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/tBGja/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:22
 [11] prepare_gradient(::typeof(f), ::AutoMooncake{Nothing}, ::Float64)
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/tBGja/src/first_order/gradient.jl:70
 [12] top-level scope
    @ REPL[43]:1
Some type information was truncated. Use `show(err)` to see complete types.

[willtebbutt]

This is a good point. I've not looked into this as extensively as I would like. It definitely ought to be straightforward though, as there's nothing particularly special about IO.

If you have any more examples of things which don't work (but ought to), please do link them!

[willtebbutt]

Ooo also, this is most certainly not the kind of error we should be getting in this setting. I'll have to make a fix for this.

[willtebbutt]

I've made #312 to fix the error that you've seen here. With it, you should see errors which look more like

┌ Warning: Unable to put rule in rule field. A `BadRuleTypeException` should be thrown.
└ @ Mooncake ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:1511
ERROR: BadRuleTypeException:

Rule is of type:
Mooncake.DerivedRule{Tuple{Base.var"##sprint#592", Nothing, Int64, typeof(sprint), Function, Tuple{Float64}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{CoDual{Base.var"##sprint#592", NoFData}, CoDual{Nothing, NoFData}, CoDual{Int64, NoFData}, CoDual{typeof(sprint), NoFData}, CoDual, CoDual{Tuple{Float64}, NoFData}}, Union{}}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{NoRData}, Tuple{NoRData, NoRData, NoRData, NoRData, Any, Tuple{Float64}}}}, Val{true}, Val{6}}

However, expected rule to be of type:
Mooncake.DerivedRule{Tuple{Base.var"##sprint#592", Nothing, Int64, typeof(sprint), Function, Tuple{Float64}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{CoDual{Base.var"##sprint#592", NoFData}, CoDual{Nothing, NoFData}, CoDual{Int64, NoFData}, CoDual{typeof(sprint), NoFData}, CoDual, CoDual{Tuple{Float64}, NoFData}}, CoDual{String, NoFData}}}, MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{NoRData}, Tuple{NoRData, NoRData, NoRData, NoRData, Any, Tuple{Float64}}}}, Val{true}, Val{6}}

This error occured for MethodInstance for Base.var"#sprint#592"(::Nothing, ::Int64, ::typeof(sprint), ::Function, ::Float64) with signature:
Tuple{Base.var"##sprint#592", Nothing, Int64, typeof(sprint), Function, Float64}

Usually this error is indicative of something having gone wrong in the compilation of the rule in question. Look at the error message for the error which caused this error (below) for more details. If the error below does not immediately give you enough information to debug what is going on, consider building the rule for the signature above, and inspecting the IR.

Stacktrace:
 [1] _build_rule!(rule::Mooncake.LazyDerivedRule{Tuple{…}, Mooncake.DerivedRule{…}}, args::Tuple{CoDual{…}, CoDual{…}, CoDual{…}, CoDual{…}, CoDual{…}, CoDual{…}})
   @ Mooncake ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:1536
 [2] LazyDerivedRule
   @ ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:1483 [inlined]
 [3] f
   @ ./REPL[10]:2 [inlined]
 [4] (::Tuple{…})(none::CoDual{…}, none::CoDual{…})
   @ Base.Experimental ./<missing>:0
 [5] (::MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{…}, CoDual{…}}})(::CoDual{typeof(f), NoFData}, ::CoDual{Float64, NoFData})
   @ MistyClosures ~/.julia/packages/MistyClosures/4rFSA/src/MistyClosures.jl:22
 [6] DerivedRule
   @ ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:739 [inlined]
 [7] __value_and_gradient!!(::Mooncake.DerivedRule{…}, ::CoDual{…}, ::CoDual{…})
   @ Mooncake ~/.julia/dev/Mooncake/src/interface.jl:73
 [8] value_and_gradient!!(::Mooncake.DerivedRule{Tuple{…}, MistyClosures.MistyClosure{…}, MistyClosures.MistyClosure{…}, Val{…}, Val{…}}, ::Function, ::Float64)
   @ Mooncake ~/.julia/dev/Mooncake/src/interface.jl:144
 [9] top-level scope
   @ REPL[15]:1

caused by: No rrule!! available for foreigncall with primal argument types Tuple{Val{:jl_alloc_string}, Val{Ref{String}}, Tuple{Val{UInt64}}, Val{1}, Val{(:ccall, 0x000e)}, UInt64}. This problem has most likely arisen because there is a ccall somewhere in the function you are trying to differentiate, for which an rrule!! has not been explicitly written.You have three options: write an rrule!! for this foreigncall, write an rrule!! for a Julia function that calls this foreigncall, or re-write your code to avoid this foreigncall entirely. If you believe that this error has arisen for some other reason than the above, or the above does not help you to workaround this problem, please open an issue.
Stacktrace:
  [1] rrule!!(::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…})
    @ Mooncake ~/.julia/dev/Mooncake/src/rrules/foreigncall.jl:12
  [2] #sprint#592
    @ ./strings/io.jl:108 [inlined]
  [3] (::Tuple{…})(none::CoDual{…}, none::CoDual{…}, none::CoDual{…}, none::CoDual{…}, none::CoDual, none::CoDual{…})
    @ Base.Experimental ./<missing>:0
  [4] (::MistyClosures.MistyClosure{…})(::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…})
    @ MistyClosures ~/.julia/packages/MistyClosures/4rFSA/src/MistyClosures.jl:22
  [5] (::Mooncake.DerivedRule{…})(::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…}, ::CoDual{…})
    @ Mooncake ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:739
  [6] _build_rule!(rule::Mooncake.LazyDerivedRule{Tuple{…}, Mooncake.DerivedRule{…}}, args::Tuple{CoDual{…}, CoDual{…}, CoDual{…}, CoDual{…}, CoDual{…}, CoDual{…}})
    @ Mooncake ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:1534
  [7] LazyDerivedRule
    @ ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:1483 [inlined]
  [8] f
    @ ./REPL[10]:2 [inlined]
  [9] (::Tuple{…})(none::CoDual{…}, none::CoDual{…})
    @ Base.Experimental ./<missing>:0
 [10] (::MistyClosures.MistyClosure{Core.OpaqueClosure{Tuple{…}, CoDual{…}}})(::CoDual{typeof(f), NoFData}, ::CoDual{Float64, NoFData})
    @ MistyClosures ~/.julia/packages/MistyClosures/4rFSA/src/MistyClosures.jl:22
 [11] DerivedRule
    @ ~/.julia/dev/Mooncake/src/interpreter/s2s_reverse_mode_ad.jl:739 [inlined]
 [12] __value_and_gradient!!(::Mooncake.DerivedRule{…}, ::CoDual{…}, ::CoDual{…})
    @ Mooncake ~/.julia/dev/Mooncake/src/interface.jl:73
 [13] value_and_gradient!!(::Mooncake.DerivedRule{Tuple{…}, MistyClosures.MistyClosure{…}, MistyClosures.MistyClosure{…}, Val{…}, Val{…}}, ::Function, ::Float64)
    @ Mooncake ~/.julia/dev/Mooncake/src/interface.jl:144
 [14] top-level scope
    @ REPL[15]:1
Some type information was truncated. Use `show(err)` to see complete types.

I think fixing this is largely going to involve the liberal use of the @zero_adjoint macro to avoid hitting ccalls associated to string manipulation. I don't have the time to look at it right now, but I would be happy to help anyone who wanted to look at it get going with it.