Enzyme error in nested jacobian call #388

vavrines · 2024-07-29T03:38:07Z

I'm having problems trying to do nested jacobian calculations. This might be related to #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

The text was updated successfully, but these errors were encountered:

gdalle · 2024-07-29T07:01:21Z

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 · 2024-07-29T08:52:43Z

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 · 2024-07-29T21:49:38Z

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

wsmoses · 2024-07-29T23:14:08Z

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 <#388 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXEU7NFXJBZCUSCULY3ZO22IRAVCNFSM6AAAAABLTNSBS6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENJXGA3DSOBZGY> . You are receiving this because you were mentioned.Message ID: ***@***.***>

gdalle · 2024-07-30T06:09:29Z

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 · 2024-08-05T15:37:41Z

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 · 2024-08-05T16:29:52Z

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 · 2024-08-08T16:39:38Z

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 · 2024-08-08T17:34:36Z

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

gdalle · 2024-09-24T12:12:21Z

This works on the current main branch of DI (probably thanks to the latest version of Enzyme which removes the need for deferred):

julia> using DifferentiationInterface, Enzyme, LinearAlgebra
Precompiling DifferentiationInterface
  2 dependencies successfully precompiled in 3 seconds. 8 already precompiled.
Precompiling DifferentiationInterfaceEnzymeExt
  1 dependency successfully precompiled in 1 seconds. 30 already precompiled.

julia> const DI = DifferentiationInterface
DifferentiationInterface

julia> X = rand(Float32, 1, 10)
1×10 Matrix{Float32}:
 0.604374  0.31886  0.28881  0.602471  0.73593  0.950794  0.891673  0.805122  0.433847  0.870221

julia> u(x) = @. x^3
u (generic function with 1 method)

julia> ux(x) = begin
           t = DI.jacobian(u, AutoEnzyme(), x)
           t |> diag |> permutedims
       end
ux (generic function with 1 method)

julia> uxx(x) = begin
           t = DI.jacobian(ux, AutoEnzyme(), x)
           t |> diag |> permutedims
       end
uxx (generic function with 1 method)

julia> u(X)
1×10 Matrix{Float32}:
 0.220759  0.0324189  0.02409  0.218679  0.398575  0.859527  0.708951  0.521897  0.0816599  0.659006

julia> ux(X)
1×10 Matrix{Float32}:
 1.09581  0.305014  0.250233  1.08891  1.62478  2.71203  2.38524  1.94466  0.564669  2.27186

julia> uxx(X)
1×10 Matrix{Float32}:
 3.62625  1.91316  1.73286  3.61482  4.41558  5.70476  5.35004  4.83073  2.60308  5.22133

gdalle added bug Something isn't working backend Related to one or more autodiff backends labels Jul 29, 2024

gdalle changed the title ~~Errored enzyme nested jacobian call~~ Enzyme error in nested jacobian call Jul 29, 2024

gdalle closed this as completed Sep 24, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enzyme error in nested jacobian call #388

Enzyme error in nested jacobian call #388

vavrines commented Jul 29, 2024

gdalle commented Jul 29, 2024

vavrines commented Jul 29, 2024

gdalle commented Jul 29, 2024

wsmoses commented Jul 29, 2024 via email

gdalle commented Jul 30, 2024

wsmoses commented Aug 5, 2024

gdalle commented Aug 5, 2024

wsmoses commented Aug 8, 2024

gdalle commented Aug 8, 2024

gdalle commented Sep 24, 2024

Enzyme error in nested jacobian call #388

Enzyme error in nested jacobian call #388

Comments

vavrines commented Jul 29, 2024

gdalle commented Jul 29, 2024

vavrines commented Jul 29, 2024

gdalle commented Jul 29, 2024

wsmoses commented Jul 29, 2024 via email

gdalle commented Jul 30, 2024

wsmoses commented Aug 5, 2024

gdalle commented Aug 5, 2024

wsmoses commented Aug 8, 2024

gdalle commented Aug 8, 2024

gdalle commented Sep 24, 2024