EnzymeAD/Enzyme.jl

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

Aberta

#1.950 aberto em 9 de out. de 2024

 (2 comentários) (0 reação) (0 responsável)Julia (96 forks)batch import
good first issue

Métricas do repositório

Stars
 (563 estrelas)
Métricas de merge de PR
 (Mesclagem média 5d 1h) (76 fundiu PRs em 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

Guia do colaborador