search

EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
444 stars 62 forks source link

Handling views properly #1950

Open gdalle opened 4 hours ago

gdalle commented 4 hours ago

As discussed with Billy in https://github.com/SciML/NonlinearSolve.jl/issues/476#issuecomment-2402984370, DI.jacobian fails when x is a SubArray. However, the reason why Enzyme succeeds is that he carefully avoided calling onehot to initialize the BatchDuplicated. Indeed, onehot creates actual Arrays because it relies on similar, and you get a type inconsistency with x.

julia> using Enzyme

julia> x = view(ones(2), 1:2)
2-element view(::Vector{Float64}, 1:2) with eltype Float64:
 1.0
 1.0

julia> Enzyme.jacobian(Enzyme.Forward, identity, x)
ERROR: MethodError: no method matching EnzymeCore.BatchDuplicated(::SubArray{Float64, 1, Vector{…}, Tuple{…}, true}, ::Tuple{Vector{…}, Vector{…}})

Closest candidates are:
  EnzymeCore.BatchDuplicated(::T1, ::Tuple{Vararg{T1, N}}) where {T1, N}
   @ EnzymeCore ~/.julia/packages/EnzymeCore/frpza/src/EnzymeCore.jl:134
  EnzymeCore.BatchDuplicated(::T1, ::Tuple{Vararg{T1, N}}, ::Bool) where {T1, N}
   @ EnzymeCore ~/.julia/packages/EnzymeCore/frpza/src/EnzymeCore.jl:134

Stacktrace:
 [1] gradient(fm::EnzymeCore.ForwardMode{…}, f::typeof(identity), x::SubArray{…}; chunk::Nothing, shadows::Tuple{…})
   @ Enzyme ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:1928
 [2] gradient
   @ ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:1909 [inlined]
 [3] #jacobian#135
   @ ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:2030 [inlined]
 [4] jacobian(::EnzymeCore.ForwardMode{…}, ::Function, ::SubArray{…})
   @ Enzyme ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:2029
 [5] top-level scope
   @ ~/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:4
Some type information was truncated. Use `show(err)` to see complete types.

I'm legitimately curious to know if there is a better solution

wsmoses commented 4 hours ago

Yeah we should specialize onehot for views here. In addition to fixing the type consistency issue, it’ll save time since we don’t need to make one hot for all inputs of the backing array, just the ones in view.

On Wed, Oct 9, 2024 at 2:03 PM Guillaume Dalle @.***> wrote:

As discussed with Billy in SciML/NonlinearSolve.jl#476 (comment) https://github.com/SciML/NonlinearSolve.jl/issues/476#issuecomment-2402984370, DI.jacobian fails when x is a SubArray. However, the reason why Enzyme succeeds is that he carefully avoided calling onehot to initialize the BatchDuplicated. Indeed, onehot creates actual Arrays because it relies on similar, and you get a type inconsistency with x.

julia> using Enzyme

julia> x = view(ones(2), 1:2)2-element view(::Vector{Float64}, 1:2) with eltype Float64: 1.0 1.0

julia> Enzyme.jacobian(Enzyme.Forward, identity, x) ERROR: MethodError: no method matching EnzymeCore.BatchDuplicated(::SubArray{Float64, 1, Vector{…}, Tuple{…}, true}, ::Tuple{Vector{…}, Vector{…}})

Closest candidates are: EnzymeCore.BatchDuplicated(::T1, ::Tuple{Vararg{T1, N}}) where {T1, N} @ EnzymeCore ~/.julia/packages/EnzymeCore/frpza/src/EnzymeCore.jl:134 EnzymeCore.BatchDuplicated(::T1, ::Tuple{Vararg{T1, N}}, ::Bool) where {T1, N} @ EnzymeCore ~/.julia/packages/EnzymeCore/frpza/src/EnzymeCore.jl:134

Stacktrace: [1] gradient(fm::EnzymeCore.ForwardMode{…}, f::typeof(identity), x::SubArray{…}; chunk::Nothing, shadows::Tuple{…}) @ Enzyme ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:1928 [2] gradient @ ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:1909 [inlined] [3] #jacobian#135 @ ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:2030 [inlined] [4] jacobian(::EnzymeCore.ForwardMode{…}, ::Function, ::SubArray{…}) @ Enzyme ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:2029 [5] top-level scope @ ~/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:4 Some type information was truncated. Use show(err) to see complete types.

I'm legitimately curious to know if there is a better solution

— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Enzyme.jl/issues/1950, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXC5YF5E7W3A3ZZKQ4TZ2V4YDAVCNFSM6AAAAABPVF7FF2VHI2DSMVQWIX3LMV43ASLTON2WKOZSGU3TMNZSGE4DQNI . You are receiving this because you are subscribed to this thread.Message ID: @.***>