[Sugar] handle view inputs for Enzyme.onehot and therefore Enzyme.gradient / Enzyme.jacobian
- 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
Assessment
This issue has not been assessed yet.