search

gdalle / DifferentiationInterface.jl

An interface to various automatic differentiation backends in Julia.
https://gdalle.github.io/DifferentiationInterface.jl/DifferentiationInterface
MIT License
163 stars 12 forks source link

Enzyme error in nested jacobian call #388

Open vavrines opened 1 month ago

vavrines commented 1 month ago

Hi @gdalle

I'm having problems trying to do nested jacobian calculations. This might be related to https://github.com/gdalle/DifferentiationInterface.jl/issues/345

MWE:

using DifferentiationInterface, Enzyme, LinearAlgebra
const DI = DifferentiationInterface

X = rand(Float32, 1, 10)

u(x) = @. x^3
ux(x) = begin
    t = DI.jacobian(u, AutoEnzyme(), x)
    t |> diag |> permutedims
end
uxx(x) = begin
    t = DI.jacobian(ux, AutoEnzyme(), x)
    t |> diag |> permutedims
end

u(X)
ux(X)
uxx(X)

Info:

ERROR: Attempting to call an indirect active function whose runtime value is inactive:
Backtrace

Stacktrace:
[1] macro expansion
  @ ~/.julia/packages/Enzyme/r8mFE/src/compiler.jl:6819
[2] enzyme_call
  @ ~/.julia/packages/Enzyme/r8mFE/src/compiler.jl:6419
[3] ForwardModeThunk
  @ ~/.julia/packages/Enzyme/r8mFE/src/compiler.jl:6299
[4] autodiff
  @ ~/.julia/packages/Enzyme/r8mFE/src/Enzyme.jl:427
[5] autodiff
  @ ~/.julia/packages/Enzyme/r8mFE/src/Enzyme.jl:326
[6] #110
  @ ~/.julia/packages/Enzyme/r8mFE/src/Enzyme.jl:1205
[7] ntuple
  @ ./ntuple.jl:19
[8] #jacobian#109
  @ ~/.julia/packages/Enzyme/r8mFE/src/Enzyme.jl:1204
gdalle commented 1 month ago

Hi, thanks for reporting this! I can reproduce the bug, but in my console there's an LLVM blurb and then it dies immediately, which makes debugging harder. Did you also encounter such behavior?

vavrines commented 1 month ago

but in my console there's an LLVM blurb and then it dies immediately

No, console is fine in my case (on a linux server).

Ver:

[a0c0ee7d] DifferentiationInterface v0.5.9
[7da242da] Enzyme v0.12.26
gdalle commented 1 month ago

@wsmoses sorry for the ping, any way you could decrypt this error message?

wsmoses commented 1 month ago

Hm, can you make a reproducer using Enzyme directly to see what’s up?

On Mon, Jul 29, 2024 at 5:50 PM Guillaume Dalle @.***> wrote:

@wsmoses https://github.com/wsmoses sorry for the ping, any way you could decrypt this error message?

— Reply to this email directly, view it on GitHub https://github.com/gdalle/DifferentiationInterface.jl/issues/388#issuecomment-2257069896, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXEU7NFXJBZCUSCULY3ZO22IRAVCNFSM6AAAAABLTNSBS6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENJXGA3DSOBZGY . You are receiving this because you were mentioned.Message ID: @.***>

gdalle commented 1 month ago

This is not exactly what DI does (it prepares the onehot and chooses a number of chunks) but it's the gist:

using Enzyme, LinearAlgebra

X = rand(Float32, 1, 10)

u(x) = @. x^3
ux(x) = begin
    t = Enzyme.jacobian(Enzyme.Forward, u, x)
    t |> diag |> permutedims
end
uxx(x) = begin
    t = Enzyme.jacobian(Enzyme.Forward, ux, x)
    t |> diag |> permutedims
end

u(X)
ux(X)
uxx(X)

Error:

julia> uxx(X)
ERROR: Attempting to call an indirect active function whose runtime value is inactive:
Backtrace

Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6658
 [2] enzyme_call
   @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6258
 [3] ForwardModeThunk
   @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6138
 [4] autodiff
   @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:427
 [5] autodiff
   @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:326
 [6] #jacobian#100
   @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:1191
 [7] #jacobian#100
   @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:0
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6658 [inlined]
  [2] enzyme_call
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6258 [inlined]
  [3] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6138 [inlined]
  [4] autodiff
    @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:427 [inlined]
  [5] autodiff
    @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:326 [inlined]
  [6] #jacobian#100
    @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:1191 [inlined]
  [7] fwddiffe10julia__jacobian_100_5017_inner_1wrap
    @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:0
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6658 [inlined]
  [9] enzyme_call
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6258 [inlined]
 [10] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6138 [inlined]
 [11] runtime_generic_fwd(activity::Type{…}, width::Val{…}, RT::Val{…}, f::Enzyme.var"##jacobian#100", df::Nothing, df_2::Nothing, df_3::Nothing, df_4::Nothing, df_5::Nothing, df_6::Nothing, df_7::Nothing, df_8::Nothing, df_9::Nothing, df_10::Nothing, primal_1::NTuple{…}, shadow_1_1::NTuple{…}, shadow_1_2::NTuple{…}, shadow_1_3::NTuple{…}, shadow_1_4::NTuple{…}, shadow_1_5::NTuple{…}, shadow_1_6::NTuple{…}, shadow_1_7::NTuple{…}, shadow_1_8::NTuple{…}, shadow_1_9::NTuple{…}, shadow_1_10::NTuple{…}, primal_2::typeof(jacobian), shadow_2_1::Nothing, shadow_2_2::Nothing, shadow_2_3::Nothing, shadow_2_4::Nothing, shadow_2_5::Nothing, shadow_2_6::Nothing, shadow_2_7::Nothing, shadow_2_8::Nothing, shadow_2_9::Nothing, shadow_2_10::Nothing, primal_3::ForwardMode{…}, shadow_3_1::Nothing, shadow_3_2::Nothing, shadow_3_3::Nothing, shadow_3_4::Nothing, shadow_3_5::Nothing, shadow_3_6::Nothing, shadow_3_7::Nothing, shadow_3_8::Nothing, shadow_3_9::Nothing, shadow_3_10::Nothing, primal_4::typeof(u), shadow_4_1::Nothing, shadow_4_2::Nothing, shadow_4_3::Nothing, shadow_4_4::Nothing, shadow_4_5::Nothing, shadow_4_6::Nothing, shadow_4_7::Nothing, shadow_4_8::Nothing, shadow_4_9::Nothing, shadow_4_10::Nothing, primal_5::Matrix{…}, shadow_5_1::Matrix{…}, shadow_5_2::Matrix{…}, shadow_5_3::Matrix{…}, shadow_5_4::Matrix{…}, shadow_5_5::Matrix{…}, shadow_5_6::Matrix{…}, shadow_5_7::Matrix{…}, shadow_5_8::Matrix{…}, shadow_5_9::Matrix{…}, shadow_5_10::Matrix{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/YDcYf/src/rules/jitrules.jl:210
 [12] jacobian
    @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:1187 [inlined]
 [13] ux
    @ ~/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:7 [inlined]
 [14] fwddiffe10julia_ux_4480wrap
    @ ~/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:0
 [15] macro expansion
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6658 [inlined]
 [16] enzyme_call
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6258 [inlined]
 [17] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/YDcYf/src/compiler.jl:6138 [inlined]
 [18] autodiff(::ForwardMode{…}, f::Const{…}, ::Type{…}, args::BatchDuplicated{…})
    @ Enzyme ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:427
 [19] autodiff
    @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:326 [inlined]
 [20] jacobian(::ForwardMode{FFIABI}, f::Function, x::Matrix{Float32}; shadow::NTuple{10, Matrix{Float32}})
    @ Enzyme ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:1191
 [21] jacobian
    @ ~/.julia/packages/Enzyme/YDcYf/src/Enzyme.jl:1187 [inlined]
 [22] uxx(x::Matrix{Float32})
    @ Main ~/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:11
 [23] top-level scope
    @ ~/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:17
Some type information was truncated. Use `show(err)` to see complete types.
wsmoses commented 1 month ago

So I think here the issue is that the inner call isn’t using deferred.

Enzyme really should auto upgrade inner calls to deferred but getting that over the line needs a bit of abstractinterpreter skill that I personally don’t have.

gdalle commented 1 month ago

Theoretically that's what DI.nested is about, but this code downright crashes my REPL before I can see what's wrong. Do you think you can help me see the Enzyme error?

using DifferentiationInterface, Enzyme, LinearAlgebra
const DI = DifferentiationInterface

X = rand(Float32, 1, 10)

u(x) = @. x^3
ux(x) = begin
    t = DI.jacobian(u, DI.nested(AutoEnzyme(mode=Enzyme.Forward)), x)
    t |> diag |> permutedims
end
uxx(x) = begin
    t = DI.jacobian(ux, AutoEnzyme(mode=Enzyme.Forward), x)
    t |> diag |> permutedims
end

u(X)
ux(X)
uxx(X)
wsmoses commented 1 month ago

Hm can you post the error message and corresponding enzyme calls again (sorry I'm not doing this to be annoying, but to see if there's something weird about how you're calling things that would trigger)

gdalle commented 1 month ago

As I said my console crashes and closes before I can see the error. Maybe it's a VScode extension thing, I'll see if it still happens outside of it