EnzymeAD / EnzymeAD/Enzyme.jl

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

Open
#1,950 2 comments 0 reactions 0 assignees View on GitHub
good first issue
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

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 `Array`s because it relies on `similar`, and you get a type inconsistency with `x`.

```julia
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

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.