EnzymeAD/Enzyme.jl

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

Ouverte

#1 950 ouverte le 9 oct. 2024

 (2 commentaires) (0 réaction) (0 personne assignée)Julia (96 forks)batch import
good first issue

Métriques du dépôt

Stars
 (563 étoiles)
Métriques de merge PR
 (Merge moyen 5j 1h) (76 PRs mergées en 30 j)

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

Guide contributeur