EnzymeAD/Enzyme.jl
[Sugar] handle view inputs for Enzyme.onehot and therefore Enzyme.gradient / Enzyme.jacobian
Open
#1,950 opened on Oct 9, 2024
good first issue
Repository metrics
- Stars
- (563 stars)
- PR merge metrics
- (Avg merge 5d 1h) (76 merged PRs in 30d)
Description
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