EnzymeAD/Enzyme.jl

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

オープン

#1,950 opened on 2024/10/09

 (2 件のコメント) (0 件のリアクション) (0 人の担当者)Julia (96 件のフォーク)batch import
good first issue

Repository metrics

Stars
 (563 個のスター)
PR merge metrics
 (平均マージ 5d 1h) (30d で 76 merged PRs)

説明

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

コントリビューターガイド