search

EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
443 stars 62 forks source link

Adding Jacobian computation support for `AbstractArray{<:AbstractArray{T}}` type function argument #1644

Closed frankwswang closed 2 months ago

frankwswang commented 2 months ago

It seems that Enzyme.jl currently does not (in default) support computing the derivatives of vector functions that take arguments of AbstractArray{<:AbstractArray{T}} as inputs.

julia> using Random

julia> Random.seed!(1234)
TaskLocalRNG()

julia> arg = [rand(2,3), rand(3,2), rand(2,2)];

julia> foo(x) = x[1] * x[2] + exp.(x[3])
foo (generic function with 1 method)

julia> foo(arg)
2×2 Matrix{Float64}:
 1.66659  2.17893
 2.65023  3.19387

julia> using Enzyme

julia> Enzyme.jacobian(Reverse, foo, arg, Val(4))
ERROR: MethodError: no method matching zero(::Type{Matrix{Float64}})
wsmoses commented 2 months ago

Yeah, in this case I think you may be better just calling autodiff directly with the corresponding one-hot shadows