[Sugar] handle view inputs for Enzyme.onehot and therefore Enzyme.gradient / Enzyme.jacobian

[gdalle]

As discussed with Billy in SciML/NonlinearSolve.jl#476 (comment), 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]

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) <SciML/NonlinearSolve.jl#476 (comment)>, 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 <#1950>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXC5YF5E7W3A3ZZKQ4TZ2V4YDAVCNFSM6AAAAABPVF7FF2VHI2DSMVQWIX3LMV43ASLTON2WKOZSGU3TMNZSGE4DQNI> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[wsmoses]

@ExpandingMan would you like to give this a go (its basically just adding the onehot function)