diff --git a/ext/EnzymeStaticArraysExt.jl b/ext/EnzymeStaticArraysExt.jl index 6dbd390cb7..bcaa3ec6cb 100644 --- a/ext/EnzymeStaticArraysExt.jl +++ b/ext/EnzymeStaticArraysExt.jl @@ -3,6 +3,8 @@ module EnzymeStaticArraysExt using StaticArrays using Enzyme +@inline Enzyme.tupstack(rows::(NTuple{N, <:StaticArrays.SArray} where N), inshape, outshape) = reshape(cat(rows..., dims=length(inshape)), (inshape..., outshape...)) + @inline function Enzyme.onehot(x::StaticArrays.SArray{S, T, N, L}) where {S, T, N, L} ntuple(Val(L)) do i Base.@_inline_meta diff --git a/src/Enzyme.jl b/src/Enzyme.jl index 583035593d..66551f2958 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -1024,12 +1024,16 @@ end end """ - gradient(::ReverseMode, f, x) + gradient(::ReverseMode, f, args...) Compute the gradient of a real-valued function `f` using reverse mode. -This will allocate and return new array `make_zero(x)` with the gradient result. +For each differentiable argument, this function will allocate and return new derivative object, returning +a tuple of derivatives for each argument. If an argument is not differentiable, the element of the returned +tuple with be nothing. -Besides arrays, for struct `x` it returns another instance of the same type, +In reverse mode (here), the derivatives will be the same type as the original argument. + +This is a structure gradient. For a struct `x` it returns another instance of the same type, whose fields contain the components of the gradient. In the result, `grad.a` contains `∂f/∂x.a` for any differential `x.a`, while `grad.c == x.c` for other types. @@ -1042,44 +1046,128 @@ f(x) = x[1]*x[2] grad = gradient(Reverse, f, [2.0, 3.0]) # output +([3.0, 2.0],) +``` -2-element Vector{Float64}: - 3.0 - 2.0 +```jldoctest gradient +grad = gradient(Reverse, only ∘ f, (a = 2.0, b = [3.0], c = "str")) + +# output + +((a = 3.0, b = [2.0], c = "str"),) ``` ```jldoctest gradient +mul(x, y) = x[1]*y[1] -grad = gradient(ReverseWithPrimal, f, [2.0, 3.0]) +grad = gradient(Reverse, mul, [2.0], [3.0]) # output -([3.0, 2.0], 6.0) +([3.0], [2.0]) ``` ```jldoctest gradient -grad = gradient(Reverse, only ∘ f, (a = 2.0, b = [3.0], c = "str")) + +grad = gradient(Reverse, mul, [2.0], Const([3.0])) + +# output +([3.0], nothing) +``` + +If passing a mode that returns the primal (e.g. ReverseWithPrimal), the return type will instead be +a tuple where the first element contains the derivatives, and the second element contains the result of the original computation. + +```jldoctest gradient + +grad = gradient(ReverseWithPrimal, f, [2.0, 3.0]) + +# output +(([3.0, 2.0],), 6.0) +``` +```jldoctest gradient + +grad = gradient(ReverseWithPrimal, mul, [2.0], [3.0]) # output +(([3.0], [2.0]), 6.0) +``` + +```jldoctest gradient +grad = gradient(ReverseWithPrimal, mul, [2.0], Const([3.0])) -(a = 3.0, b = [2.0], c = "str") +# output +(([3.0], nothing), 6.0) ``` + """ -@inline function gradient(rm::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}, f::F, x::X) where {F, X, ReturnPrimal, RuntimeActivity, ABI, Holomorphic, ErrIfFuncWritten} - if Compiler.active_reg_inner(X, #=seen=#(), #=world=#nothing, #=justActive=#Val(true)) == Compiler.ActiveState - dx = Ref(make_zero(x)) - res = autodiff(rm, f, Active, MixedDuplicated(x, dx)) - if ReturnPrimal - (only(dx), res[2]) - else - only(dx) +@generated function gradient(rm::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}, f::F, x::ty_0, args::Vararg{<:Any, N}) where {F, ty_0, ReturnPrimal, RuntimeActivity, ABI, Holomorphic, ErrIfFuncWritten, N} + toemit= Expr[quote + act_0 = !(x isa Enzyme.Const) && Compiler.active_reg_inner(Core.Typeof(x), #=seen=#(), #=world=#nothing, #=justActive=#Val(true)) == Compiler.ActiveState + end] + rargs = Union{Symbol,Expr}[:x] + acts = Symbol[Symbol("act_0")] + + for i in 1:N + argidx = quote args[$i] end + push!(rargs, argidx) + sym = Symbol("act_$i") + push!(acts, sym) + push!(toemit, quote + $sym = !($argidx isa Enzyme.Const) && Compiler.active_reg_inner(Core.Typeof($argidx), #=seen=#(), #=world=#nothing, #=justActive=#Val(true)) == Compiler.ActiveState + end) + end + + idx = 0 + shadows = Symbol[] + enz_args = Expr[] + resargs = Expr[] + for (arg, act) in zip(rargs, acts) + shad = Symbol("shad_$idx") + push!(shadows, shad) + push!(toemit, quote + $shad = if $arg isa Enzyme.Const + nothing + elseif $act + Ref(make_zero($arg)) + else + make_zero($arg) + end + end) + push!(enz_args, quote + if $arg isa Enzyme.Const + $arg + elseif $act + MixedDuplicated($arg, $shad) + else + Duplicated($arg, $shad) + end + end) + push!(resargs, quote + if $arg isa Enzyme.Const + nothing + elseif $act + $shad[] + else + $shad + end + end) + idx+=1 + end + push!(toemit, quote + res = autodiff(rm, f, Active, $(enz_args...)) + end) + + if ReturnPrimal + return quote + Base.@_inline_meta + $(toemit...) + (($(resargs...),), res[2]) end else - dx = make_zero(x) - res = autodiff(rm, f, Active, Duplicated(x, dx)) - if ReturnPrimal - (dx, res[2]) - else - dx + return quote + Base.@_inline_meta + $(toemit...) + ($(resargs...),) end end end @@ -1100,10 +1188,7 @@ dx = [0.0, 0.0] gradient!(Reverse, dx, f, [2.0, 3.0]) # output - -2-element Vector{Float64}: - 3.0 - 2.0 +([3.0, 2.0],) ``` ```jldoctest gradip @@ -1111,21 +1196,87 @@ dx = [0.0, 0.0] gradient!(ReverseWithPrimal, dx, f, [2.0, 3.0]) # output -([3.0, 2.0], 6.0) +(([3.0, 2.0],), 6.0) ``` """ @inline function gradient!(rm::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}, dx::X, f::F, x::X) where {X<:Array, F, ReturnPrimal, RuntimeActivity, ABI, Holomorphic, ErrIfFuncWritten} make_zero!(dx) res = autodiff(rm, f, Active, Duplicated(x, dx)) return if ReturnPrimal - (dx, res[2]) + ((dx,), res[2]) else - dx + (dx,) + end +end + +@inline function chunkedonehot(x, ::Val{chunk}) where chunk + sz = length(x) + num = ((sz + chunk - 1) ÷ chunk) + ntuple(Val(num)) do i + Base.@_inline_meta + onehot(x, (i-1)*chunk+1, i == num ? sz : (i*chunk) ) + end +end + +@inline function chunkedonehot(x::AbstractFloat, ::Val{chunk}) where chunk + return ((one(x),),) +end + +@inline tupleconcat(x) = x +@inline tupleconcat(x, y) = (x..., y...) +@inline tupleconcat(x, y, z...) = (x..., tupleconcat(y, z...)...) + +function create_shadows(::Nothing, x) + return (onehot(x),) +end + +function create_shadows(::Val{1}, x) + return (onehot(x),) +end + +function create_shadows(::Val{chunk}, x) where chunk + return (chunkedonehot(x, Val(chunk)),) +end + +struct TupleArray{T, Shape, Length, N} <: AbstractArray{T,N} + data::NTuple{Length, T} +end +TupleArray(data::NTuple{Length, T}, Shape) where {Length, T} = TupleArray{T, Shape, Length, length(Shape)}(data) + +@inline Base.eltype(::TupleArray{T}) where T = T +@inline Base.eltype(::Type{<:TupleArray{T}}) where T = T +@inline Base.size(::TupleArray{<:Any, Shape}) where Shape = Shape +@inline Base.ndims(::TupleArray{<:Any, <:Any, <:Any, N}) where N = N + +function Base.convert(::Type{Array{T, N}}, X::TupleArray{T, Shape, Length, N}) where {T, Shape, Length, N} + vals = Array{T, N}(undef, Shape...) + for i in 1:Length + @inbounds val[i] = X.data[i] + end + return vals +end + +function Base.getindex(a::TupleArray, args::Vararg{Int,N}) where {N} + start = 0 + for i in 1:N + start *= size(a, N - i + 1) + start += (args[N - i + 1] - 1) + end + start += 1 + return a.data[start] +end + +@inline function tupstack(x, inshape, outshape) + st = Base.stack(x) + if length(outshape) == 1 + st + else + reshape(st, (inshape..., outshape...)) end end """ - gradient(::ForwardMode, f, x; shadow=onehot(x)) + gradient(::ForwardMode, f, x; shadows=onehot(x), chunk=nothing) Compute the gradient of an array-input function `f` using forward mode. The optional keyword argument `shadow` is a vector of one-hot vectors of type `x` @@ -1138,372 +1289,331 @@ Example: ```jldoctest gradfwd f(x) = x[1]*x[2] -grad = gradient(Forward, f, [2.0, 3.0]) +gradient(Forward, f, [2.0, 3.0]) # output -(3.0, 2.0) +([3.0, 2.0],) ``` ```jldoctest gradfwd gradient(ForwardWithPrimal, f, [2.0, 3.0]) # output -((3.0, 2.0), 6.0) +(([3.0, 2.0],), 6.0) ``` -""" -@inline function gradient(fm::ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity}, f, x; shadow=onehot(x)) where {ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity} - if length(shadow) == 0 - if ReturnPrimal - ((), f(x.val)) - else - return () - end - end - resp = autodiff(fm, f, BatchDuplicated, BatchDuplicated(x, shadow)) - res = values(resp[1]) - dres = if x isa AbstractFloat - res[1] - else - res - end - if ReturnPrimal - (dres, resp[2]) - else - dres - end -end - -@inline function chunkedonehot(x, ::Val{chunk}) where chunk - sz = length(x) - num = ((sz + chunk - 1) ÷ chunk) - ntuple(Val(num)) do i - Base.@_inline_meta - onehot(x, (i-1)*chunk+1, i == num ? sz : (i*chunk) ) - end -end +```jldoctest gradfwd +gradient(Forward, f, [2.0, 3.0]; chunk=Val(1)) -@inline function chunkedonehot(x::AbstractFloat, ::Val{chunk}) where chunk - return ((one(x),),) -end +# output -@inline tupleconcat(x) = x -@inline tupleconcat(x, y) = (x..., y...) -@inline tupleconcat(x, y, z...) = (x..., tupleconcat(y, z...)...) +([3.0, 2.0],) +``` -""" - gradient(::ForwardMode, f, x::Union{Array,NTuple}, ::Val{chunk}; shadow=onehot(x)) +```jldoctest gradfwd +gradient(ForwardWithPrimal, f, [2.0, 3.0]; chunk=Val(1)) -Compute the gradient of an array-input function `f` using vector forward mode. -Like [`gradient`](@ref), except it uses a chunk size of `chunk` to compute -`chunk` derivatives in a single call. +# output +(([3.0, 2.0],), 6.0) +``` -Example: +For functions which return an AbstractArray or scalar, this function will return an AbstracttArray +whose shape is `(size(output)..., size(input)...)`. No guarantees are presently made +about the type of the AbstractArray returned by this function (which may or may not be the same +as the input AbstractArray if provided). +For functions who return other types, this function will retun an AbstractArray +of shape `size(input)` of values of the output type. ```jldoctest -f(x) = x[1]*x[2] +f(x) = [ x[1] * x[2], x[2] + x[3] ] -grad = gradient(Forward, f, [2.0, 3.0], Val(2)) +grad = gradient(Forward, f, [2.0, 3.0, 4.0]) # output - -(3.0, 2.0) +([3.0 2.0 0.0; 0.0 1.0 1.0],) ``` """ -@inline function gradient(fm::ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity}, f::F, x::X, ::Val{chunk}; shadow=chunkedonehot(x, Val(chunk))) where {F, X, chunk, ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity} - if chunk == 0 - throw(ErrorException("Cannot differentiate with a batch size of 0")) - end - if ReturnPrimal - rp = autodiff(fm, f, BatchDuplicated, BatchDuplicated(x, shadow[1]))[1] - dres1 = if chunk == 1 - (rp[1],) - else - values(rp[1]) - end - gres = if x isa AbstractFloat - dres1 +@inline function gradient(fm::ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity}, f, x; chunk::CS=nothing, shadows=create_shadows(chunk, x)) where {ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity, CS} + if length(shadows[1]) == 0 + if ReturnPrimal + ((x,), f(x.val)) else - fm2 = ForwardMode{#=ReturnPrimal=#false, ABI, ErrIfFuncWritten,RuntimeActivity}() - tmp = ntuple(length(shadow)-1) do i - values(autodiff(fm2, f, BatchDuplicated, BatchDuplicated(x, shadow[i+1]))[1]) - end - tupleconcat(dres1, tmp...) - end - (gres, rp[2]) - else - tmp = ntuple(length(shadow)) do i - values(autodiff(fm, f, BatchDuplicated, BatchDuplicated(x, shadow[i]))[1]) + return (x,) end - res = tupleconcat(tmp...) - if x isa AbstractFloat + end + if chunk == Val(0) + throw(ErrorException("Cannot differentiate with a batch size of 0")) + end + + gradtup = if chunk == nothing + resp = autodiff(fm, f, BatchDuplicated, BatchDuplicated(x, shadows[1])) + + res = values(resp[1]) + dres = if x isa AbstractFloat res[1] else res end - end -end - -@inline function gradient(fm::ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity}, f::F, x::X, ::Val{1}; shadow=onehot(x)) where {F, X, ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity} - if ReturnPrimal - rp = autodiff(fm, f, Duplicated, Duplicated(x, shadow[1])) - dres1 = rp[1] - fm2 = ForwardMode{#=ReturnPrimal=#false, ABI, ErrIfFuncWritten,RuntimeActivity}() - - res = ntuple(length(shadow)-1) do i - autodiff(fm2, f, Duplicated, Duplicated(x, shadow[i+1]))[1] + if ReturnPrimal + ((dres,), resp[2]) + else + (dres,) end - gres = if x isa AbstractFloat - dres1 + elseif chunk == Val(1) + if ReturnPrimal + rp = autodiff(fm, f, Duplicated, Duplicated(x, shadows[1][1])) + dres1 = rp[1] + fm2 = ForwardMode{#=ReturnPrimal=#false, ABI, ErrIfFuncWritten,RuntimeActivity}() + + res = ntuple(length(shadows[1])-1) do i + autodiff(fm2, f, Duplicated, Duplicated(x, shadows[1][i+1]))[1] + end + gres = if x isa AbstractFloat + dres1[1] + else + (dres1, res...) + end + ((gres,), rp[2]) else - (dres1, res...) + res = ntuple(length(shadows[1])) do i + autodiff(fm, f, Duplicated, Duplicated(x, shadows[1][i]))[1] + end + (if x isa AbstractFloat + res[1] + else + res + end,) end - (gres, rp[2]) else - res = ntuple(length(shadow)) do i - autodiff(fm, f, Duplicated, Duplicated(x, shadow[i]))[1] - end - if x isa AbstractFloat - res[1] + if ReturnPrimal + rp = autodiff(fm, f, BatchDuplicated, BatchDuplicated(x, shadows[1][1])) + dres1 = values(rp[1]) + gres = if x isa AbstractFloat + dres1[1] + else + fm2 = ForwardMode{#=ReturnPrimal=#false, ABI, ErrIfFuncWritten,RuntimeActivity}() + tmp = ntuple(length(shadows[1])-1) do i + values(autodiff(fm2, f, BatchDuplicated, BatchDuplicated(x, shadows[1][i+1]))[1]) + end + tupleconcat(dres1, tmp...) + end + ((gres,), rp[2]) else - res + tmp = ntuple(length(shadows[1])) do i + values(autodiff(fm, f, BatchDuplicated, BatchDuplicated(x, shadows[1][i]))[1]) + end + res = tupleconcat(tmp...) + (if x isa AbstractFloat + res[1] + else + res + end,) end end -end - -""" - jacobian(::ForwardMode, f, x; shadow=onehot(x)) - jacobian(::ForwardMode, f, x, ::Val{chunk}; shadow=onehot(x)) - -Compute the jacobian of an array or scalar-input function `f` using (potentially vector) -forward mode. All relevant arguments of the forward-mode [`gradient`](@ref) function -apply here. - -Example: - -```jldoctest -f(x) = [ x[1] * x[2], x[2] + x[3] ] - -grad = jacobian(Forward, f, [2.0, 3.0, 4.0]) - -# output -2×3 Matrix{Float64}: - 3.0 2.0 0.0 - 0.0 1.0 1.0 -``` - -For functions which return an AbstractArray, this function will return an array -whose shape is `(size(output)..., size(input)...)` - -For functions who return other types, this function will retun an array or tuple -of shape `size(input)` of values of the output type. -""" -@inline function jacobian(fm::ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity}, args...; kwargs...) where {ReturnPrimal, ABI, ErrIfFuncWritten,RuntimeActivity} - gradtup = gradient(fm, args...; kwargs...) cols = if ReturnPrimal - gradtup[1] + gradtup[1][1] else - gradtup + gradtup[1] end - x = args[2] res = if x isa AbstractFloat cols - elseif length(cols) > 0 && cols[1] isa AbstractArray + elseif length(cols) > 0 && cols[1] isa AbstractArray && x isa AbstractArray inshape = size(x) outshape = size(cols[1]) # st : outshape x total inputs - st = Base.stack(cols) - - st3 = if length(inshape) <= 1 - st - else - reshape(st, (outshape..., inshape...)) - end - - st3 + tupstack(cols, outshape, inshape) elseif x isa AbstractArray - inshape = size(x) - reshape(collect(cols), inshape) + TupleArray(cols, size(x)) else cols end if ReturnPrimal - (res, gradtup[2]) + ((res,), gradtup[2]) else - res + (res,) end end """ - jacobian(::ReverseMode, f, x, ::Val{num_outs}, ::Val{chunk}=Val(1)) + jacobian(::ForwardMode, args...; kwargs...) + +Equivalent to gradient(::ForwardMode, args...; kwargs...) +""" +@inline function jacobian(fm::ForwardMode, args...; kwargs...) + gradient(fm, args...; kwargs...) +end + +""" + jacobian(::ReverseMode, f, x; n_outs=nothing, chunk=nothing) jacobian(::ReverseMode, f, x) -Compute the jacobian of an array-output function `f` using (potentially vector) -reverse mode. The `chunk` argument denotes the chunk size to use and `num_outs` -denotes the number of outputs `f` will return in an array. +Compute the jacobian of a array-output function `f` using (potentially vector) +reverse mode. The `chunk` argument denotes the chunk size to use and `n_outs` +denotes the shape of the array returned by `f`. Example: ```jldoctest f(x) = [ x[1] * x[2], x[2] + x[3] ] -grad = jacobian(Reverse, f, [2.0, 3.0, 4.0], Val(2)) +jacobian(Reverse, f, [2.0, 3.0, 4.0]) # output +([3.0 2.0 0.0; 0.0 1.0 1.0],) +``` -2×3 transpose(::Matrix{Float64}) with eltype Float64: - 3.0 2.0 0.0 - 0.0 1.0 1.0 +```jldoctest +f(x) = [ x[1] * x[2], x[2] + x[3] ] + +grad = jacobian(Reverse, f, [2.0, 3.0, 4.0], n_outs=Val((2,))) + +# output +([3.0 2.0 0.0; 0.0 1.0 1.0],) ``` -For functions which return an AbstractArray, this function will return an array -whose shape is `(size(output)..., size(input)...)` +This function will return an AbstractArray whose shape is `(size(output)..., size(input)...)`. +No guarantees are presently made about the type of the AbstractArray returned by this function +(which may or may not be the same as the input AbstractArray if provided). -For functions who return other types, this function will retun an array or tuple -of shape `size(output)` of values of the input type. +In the future, when this function is extended to handle non-array return types, +this function will retun an AbstractArray of shape `size(output)` of values of the input type. ``` """ -@inline function jacobian(::ReverseMode{#=ReturnPrimal=#false,RuntimeActivity, RABI, #=Holomorphic=#false, ErrIfFuncWritten}, f::F, x::X, n_outs::Val{n_out_val}, ::Val{chunk}) where {F, X, chunk, n_out_val, RABI<:ABI, ErrIfFuncWritten, RuntimeActivity} - num = ((n_out_val + chunk - 1) ÷ chunk) - - if chunk == 0 - throw(ErrorException("Cannot differentiate with a batch size of 0")) - end - - XT = Core.Typeof(x) - MD = Compiler.active_reg_inner(XT, #=seen=#(), #=world=#nothing, #=justActive=#Val(true)) == Compiler.ActiveState - tt′ = MD ? Tuple{BatchMixedDuplicated{XT, chunk}} : Tuple{BatchDuplicated{XT, chunk}} - tt = Tuple{XT} - rt = Core.Compiler.return_type(f, tt) - ModifiedBetween = Val((false, false)) - FA = Const{Core.Typeof(f)} - opt_mi = if RABI <: NonGenABI - Compiler.fspec(eltype(FA), tt′) - else - Val(codegen_world_age(Core.Typeof(f), tt)) - end - primal, adjoint = Enzyme.Compiler.thunk(opt_mi, FA, BatchDuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(chunk), ModifiedBetween, #=ReturnPrimal=#Val(false), #=ShadowInit=#Val(false), RABI, Val(ErrIfFuncWritten), Val(RuntimeActivity)) - - if num * chunk == n_out_val - last_size = chunk - primal2, adjoint2 = primal, adjoint - else - last_size = n_out_val - (num-1)*chunk - tt′ = Tuple{BatchDuplicated{Core.Typeof(x), last_size}} - primal2, adjoint2 = Enzyme.Compiler.thunk(opt_mi, FA, BatchDuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(last_size), ModifiedBetween, #=ReturnPrimal=#Val(false), #=ShadowInit=#Val(false), RABI, Val(ErrIfFuncWritten), Val(RuntimeActivity)) - end +@inline function jacobian(::ReverseMode{ReturnPrimal,RuntimeActivity, RABI, Holomorphic, ErrIfFuncWritten}, f::F, x::X; n_outs::OutType=nothing, chunk::CT=nothing) where {ReturnPrimal, F, X, RABI<:ABI, ErrIfFuncWritten, RuntimeActivity, OutType, CT, Holomorphic} - tmp = ntuple(num) do i - Base.@_inline_meta - dx = ntuple(Val(i == num ? last_size : chunk)) do idx - Base.@_inline_meta - z = make_zero(x) - MD ? Ref(z) : z - end - res = (i == num ? primal2 : primal)(Const(f), MD ? BatchMixedDuplicated(x, dx) : BatchDuplicated(x, dx)) - tape = res[1] - j = 0 - for shadow in res[3] - j += 1 - @inbounds shadow[(i-1)*chunk+j] += Compiler.default_adjoint(eltype(typeof(shadow))) + if n_outs == nothing + res = if f isa Const + f.val(x) + else + f(x) end - (i == num ? adjoint2 : adjoint)(Const(f), MD ? BatchMixedDuplicated(x, dx) : BatchDuplicated(x, dx), tape) - return MD ? (ntuple(Val(i == num ? last_size : chunk)) do idx - Base.@_inline_meta - dx[idx][] - end) : dx, (i == 1 ? size(res[3][1]) : nothing) - end - rows = tupleconcat(map(first, tmp)...) - outshape = tmp[1][2] - if x isa AbstractArray - inshape = size(x) - - st = Base.stack(rows) - - st2 = if length(outshape) == 1 - st + jac = if res isa AbstractArray + jacobian(ReverseMode{false,RuntimeActivity,RABI, Holomorphic, ErrIfFuncWritten}(), f, x; n_outs=Val(size(res)), chunk) + elseif res isa AbstractFloat + gradient(ReverseMode{false,RuntimeActivity,RABI, Holomorphic, ErrIfFuncWritten}(), f, x) else - reshape(st, (inshape..., outshape...)) + throw(AssertionError("Unsupported return type of function for reverse-mode jacobian, $(Core.Typeof(res))")) end - st3 = if length(outshape) == 1 && length(inshape) == 1 - transpose(st2) + return if ReturnPrimal + (jac, res) else - transp = ( ((length(inshape)+1):(length(inshape)+length(outshape)))... , (1:length(inshape))... ) - PermutedDimsArray(st2, transp) + jac end - - st3 else - reshape(collect(rows), outshape) - end -end - -@inline function jacobian(::ReverseMode{#=ReturnPrimal=#false,RuntimeActivity,RABI, #=Holomorphic=#false, ErrIfFuncWritten}, f::F, x::X, n_outs::Val{n_out_val}, ::Val{1} = Val(1)) where {F, X, n_out_val,RuntimeActivity,RABI<:ABI, ErrIfFuncWritten} - XT = Core.Typeof(x) - MD = Compiler.active_reg_inner(XT, #=seen=#(), #=world=#nothing, #=justActive=#Val(true)) == Compiler.ActiveState - tt′ = MD ? Tuple{MixedDuplicated{XT}} : Tuple{Duplicated{XT}} - tt = Tuple{XT} - rt = Core.Compiler.return_type(f, tt) - ModifiedBetween = Val((false, false)) - FA = Const{Core.Typeof(f)} - opt_mi = if RABI <: NonGenABI - Compiler.fspec(eltype(FA), tt′) - else - Val(codegen_world_age(Core.Typeof(f), tt)) - end - primal, adjoint = Enzyme.Compiler.thunk(opt_mi, FA, DuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(1), ModifiedBetween, #=ReturnPrimal=#Val(false), #=ShadowInit=#Val(false), RABI, Val(ErrIfFuncWritten), Val(RuntimeActivity)) - tmp = ntuple(n_outs) do i - Base.@_inline_meta - z = make_zero(x) - dx = MD ? Ref(z) : z - res = primal(Const(f), MD ? MixedDuplicated(x, dx) : Duplicated(x, dx)) - tape = res[1] - @inbounds res[3][i] += Compiler.default_adjoint(eltype(typeof(res[3]))) - adjoint(Const(f), MD ? MixedDuplicated(x, dx) : Duplicated(x, dx), tape) - return MD ? dx[] : dx, (i == 1 ? size(res[3]) : nothing) - end - rows = map(first, tmp) - outshape = tmp[1][2] - if x isa AbstractArray - inshape = size(x) - st = Base.stack(rows) - - st2 = if length(outshape) == 1 - st + @assert !Holomorphic + n_out_val = if length(Compiler.element(n_outs)) == 0 + 0 else - reshape(st, (inshape..., outshape...)) + prod(Compiler.element(n_outs)) end - - st3 = if length(outshape) == 1 && length(inshape) == 1 - transpose(st2) + + if chunk == Val(0) + throw(ErrorException("Cannot differentiate with a batch size of 0")) + end + + XT = Core.Typeof(x) + MD = Compiler.active_reg_inner(XT, #=seen=#(), #=world=#nothing, #=justActive=#Val(true)) == Compiler.ActiveState + tt = Tuple{XT} + rt = if f isa Const + Core.Compiler.return_type(f.val, tt) + else + Core.Compiler.return_type(f, tt) + end + + ModifiedBetween = Val((false, false)) + FRT = Core.Typeof(f) + FA = Const{FRT} + + opt_mi = if RABI <: NonGenABI + Compiler.fspec(FRT, tt′) else - transp = ( ((length(inshape)+1):(length(inshape)+length(outshape)))... , (1:length(inshape))... ) - PermutedDimsArray(st2, transp) + Val(codegen_world_age(FRT, tt)) end - st3 - else - reshape(collect(rows), outshape) - end -end + if chunk == Val(1) || chunk == nothing + tt′ = MD ? Tuple{MixedDuplicated{XT}} : Tuple{Duplicated{XT}} + primal, adjoint = Enzyme.Compiler.thunk(opt_mi, FA, DuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(1), ModifiedBetween, #=ReturnPrimal=#Val(false), #=ShadowInit=#Val(false), RABI, Val(ErrIfFuncWritten), Val(RuntimeActivity)) + tmp = ntuple(Val(n_out_val)) do i + Base.@_inline_meta + z = make_zero(x) + dx = MD ? Ref(z) : z + res = primal(Const(f), MD ? MixedDuplicated(x, dx) : Duplicated(x, dx)) + tape = res[1] + @inbounds res[3][i] += Compiler.default_adjoint(eltype(typeof(res[3]))) + adjoint(Const(f), MD ? MixedDuplicated(x, dx) : Duplicated(x, dx), tape) + return MD ? dx[] : dx, (i == 1 ? size(res[3]) : nothing) + end + rows = map(first, tmp) + outshape = tmp[1][2] + rows, outshape + else + chunksize = Compiler.element(chunk) + tt′ = MD ? Tuple{BatchMixedDuplicated{XT, chunksize}} : Tuple{BatchDuplicated{XT, chunksize}} + primal, adjoint = Enzyme.Compiler.thunk(opt_mi, FA, BatchDuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#chunk, ModifiedBetween, #=ReturnPrimal=#Val(false), #=ShadowInit=#Val(false), RABI, Val(ErrIfFuncWritten), Val(RuntimeActivity)) + + num = ((n_out_val + chunksize - 1) ÷ chunksize) + + if num * chunksize == n_out_val + last_size = chunksize + primal2, adjoint2 = primal, adjoint + else + last_size = n_out_val - (num-1)*chunksize + tt′ = Tuple{BatchDuplicated{Core.Typeof(x), last_size}} + primal2, adjoint2 = Enzyme.Compiler.thunk(opt_mi, FA, BatchDuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(last_size), ModifiedBetween, #=ReturnPrimal=#Val(false), #=ShadowInit=#Val(false), RABI, Val(ErrIfFuncWritten), Val(RuntimeActivity)) + end -@inline function jacobian(::ReverseMode{ReturnPrimal,RuntimeActivity, RABI, Holomorphic, ErrIfFuncWritten}, f::F, x::X) where {ReturnPrimal, F, X, RABI<:ABI, ErrIfFuncWritten, RuntimeActivity, Holomorphic} - res = f(x) - jac = if res isa AbstractArray - jacobian(ReverseMode{false,RuntimeActivity,RABI, Holomorphic, ErrIfFuncWritten}(), f, x, Val(length(res))) - elseif res isa AbstractFloat - gradient(ReverseMode{false,RuntimeActivity,RABI, Holomorphic, ErrIfFuncWritten}(), f, x) - else - throw(AssertionError("Unsupported return type of function for reverse-mode jacobian, $(Core.Typeof(res))")) - end + tmp = ntuple(num) do i + Base.@_inline_meta + dx = ntuple(Val(i == num ? last_size : chunksize)) do idx + Base.@_inline_meta + z = make_zero(x) + MD ? Ref(z) : z + end + res = (i == num ? primal2 : primal)(Const(f), MD ? BatchMixedDuplicated(x, dx) : BatchDuplicated(x, dx)) + tape = res[1] + j = 0 + for shadow in res[3] + j += 1 + @inbounds shadow[(i-1)*chunksize+j] += Compiler.default_adjoint(eltype(typeof(shadow))) + end + (i == num ? adjoint2 : adjoint)(Const(f), MD ? BatchMixedDuplicated(x, dx) : BatchDuplicated(x, dx), tape) + return MD ? (ntuple(Val(i == num ? last_size : chunksize)) do idx + Base.@_inline_meta + dx[idx][] + end) : dx, (i == 1 ? size(res[3][1]) : nothing) + end + rows = tupleconcat(map(first, tmp)...) + outshape = tmp[1][2] + rows, outshape + end + res = if x isa AbstractArray + inshape = size(x) + st2 = tupstack(rows, inshape, outshape) - if ReturnPrimal - (res, jac) - else - jac + st3 = if length(outshape) == 1 && length(inshape) == 1 + transpose(st2) + else + transp = ( ((length(inshape)+1):(length(inshape)+length(outshape)))... , (1:length(inshape))... ) + PermutedDimsArray(st2, transp) + end + + st3 + else + reshape(collect(rows), outshape) + end + if ReturnPrimal + # TODO optimize away redundant fwd pass + (res, if f isa Enzyme.Const + f.val(x) + else + f(x) + end) + else + (res,) + end end end diff --git a/test/ext/logexpfunctions.jl b/test/ext/logexpfunctions.jl index 69ee7f2e73..51dbe2ec76 100644 --- a/test/ext/logexpfunctions.jl +++ b/test/ext/logexpfunctions.jl @@ -9,6 +9,6 @@ xlogydiff(x) = xlogy(x[1], 23.0) grad_forward = Enzyme.gradient(Enzyme.Forward, xlogydiff, x) grad_reverse = Enzyme.gradient(Enzyme.Reverse, xlogydiff, x) - @test grad_forward[1] ≈ log(23.0) - @test grad_reverse[1] ≈ log(23.0) + @test grad_forward[1] ≈ [log(23.0)] + @test grad_reverse[1] ≈ [log(23.0)] end diff --git a/test/internal_rules.jl b/test/internal_rules.jl index b9a705941c..32a206c62e 100644 --- a/test/internal_rules.jl +++ b/test/internal_rules.jl @@ -197,14 +197,14 @@ end dL = zero(x) dL[2, 1] = 1.0 - @test Enzyme.gradient(Reverse, chol_lower0, x) ≈ [0.05270807565639164 0.0 0.0 0.0; 1.0000000000000024 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] + @test Enzyme.gradient(Reverse, chol_lower0, x)[1] ≈ [0.05270807565639164 0.0 0.0 0.0; 1.0000000000000024 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] - @test reshape(collect(Enzyme.gradient(Forward, chol_lower0, x)), 4, 4) ≈ [0.05270807565639164 0.0 0.0 0.0; 1.0000000000000024 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] + @test Enzyme.gradient(Forward, chol_lower0, x)[1] ≈ [0.05270807565639164 0.0 0.0 0.0; 1.0000000000000024 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] @test FiniteDifferences.grad(central_fdm(5, 1), chol_lower0, x)[1] ≈ [0.05270807565639164 0.0 0.0 0.0; 1.0000000000000024 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] - @test reshape(collect(Enzyme.gradient(Forward, chol_upper0, x)), 4, 4) ≈ [0.05270807565639728 0.9999999999999999 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] - @test Enzyme.gradient(Reverse, chol_upper0, x) ≈ [0.05270807565639728 0.9999999999999999 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] + @test Enzyme.gradient(Forward, chol_upper0, x)[1] ≈ [0.05270807565639728 0.9999999999999999 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] + @test Enzyme.gradient(Reverse, chol_upper0, x)[1] ≈ [0.05270807565639728 0.9999999999999999 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] @test FiniteDifferences.grad(central_fdm(5, 1), chol_upper0, x)[1] ≈ [0.05270807565639728 0.9999999999999999 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] end @@ -225,14 +225,14 @@ end x = [1.0 0.13147601759884564 0.5282944836504488; 0.13147601759884564 1.0 0.18506733179093515; 0.5282944836504488 0.18506733179093515 1.0] for i in 1:size(x, 1) for j in 1:size(x, 2) - reverse_grad = Enzyme.gradient(Reverse, x -> tchol_lower(x, i, j), x) - forward_grad = reshape(collect(Enzyme.gradient(Forward, x -> tchol_lower(x, i, j), x)), size(x)) + reverse_grad = Enzyme.gradient(Reverse, x -> tchol_lower(x, i, j), x)[1] + forward_grad = Enzyme.gradient(Forward, x -> tchol_lower(x, i, j), x)[1] finite_diff = FiniteDifferences.grad(central_fdm(5, 1), x -> tchol_lower(x, i, j), x)[1] @test reverse_grad ≈ finite_diff @test forward_grad ≈ finite_diff - reverse_grad = Enzyme.gradient(Reverse, x -> tchol_upper(x, i, j), x) - forward_grad = reshape(collect(Enzyme.gradient(Forward, x -> tchol_upper(x, i, j), x)), size(x)) + reverse_grad = Enzyme.gradient(Reverse, x -> tchol_upper(x, i, j), x)[1] + forward_grad = Enzyme.gradient(Forward, x -> tchol_upper(x, i, j), x)[1] finite_diff = FiniteDifferences.grad(central_fdm(5, 1), x -> tchol_upper(x, i, j), x)[1] @test reverse_grad ≈ finite_diff @test forward_grad ≈ finite_diff @@ -257,26 +257,26 @@ end x = [1.0 0.13147601759884564 0.5282944836504488; 0.13147601759884564 1.0 0.18506733179093515; 0.5282944836504488 0.18506733179093515 1.0] for i in 1:15 B = [3.1 2.7 5.9 2.4 1.6; 7.9 8.2 1.3 9.4 5.5; 4.7 2.9 9.8 7.1 4.3] - reverse_grad = Enzyme.gradient(Reverse, Const(B -> tcholsolv_lower(x, B, i)), B) - # forward_grad = reshape(collect(Enzyme.gradient(Forward, B -> tcholsolv_lower(x, B, i), B)), size(B)) + reverse_grad = Enzyme.gradient(Reverse, Const(B -> tcholsolv_lower(x, B, i)), B)[1] + # forward_grad = Enzyme.gradient(Forward, B -> tcholsolv_lower(x, B, i), B)[1] finite_diff = FiniteDifferences.grad(central_fdm(5, 1), B -> tcholsolv_lower(x, B, i), B)[1] @test reverse_grad ≈ finite_diff # @test forward_grad ≈ finite_diff - reverse_grad = Enzyme.gradient(Reverse, Const(B -> tcholsolv_upper(x, B, i)), B) - # forward_grad = reshape(collect(Enzyme.gradient(Forward, B -> tcholsolv_upper(x, B, i), B)), size(B)) + reverse_grad = Enzyme.gradient(Reverse, Const(B -> tcholsolv_upper(x, B, i)), B)[1] + # forward_grad = Enzyme.gradient(Forward, B -> tcholsolv_upper(x, B, i), B))[1] finite_diff = FiniteDifferences.grad(central_fdm(5, 1), B -> tcholsolv_upper(x, B, i), B)[1] @test reverse_grad ≈ finite_diff # @test forward_grad ≈ finite_diff - reverse_grad = Enzyme.gradient(Reverse, Const(x -> tcholsolv_lower(x, B, i)), x) - #forward_grad = reshape(collect(Enzyme.gradient(Forward, x -> tcholsolv_lower(x, B, i), x)), size(x)) + reverse_grad = Enzyme.gradient(Reverse, Const(x -> tcholsolv_lower(x, B, i)), x)[1] + #forward_grad = Enzyme.gradient(Forward, x -> tcholsolv_lower(x, B, i), x)[1] finite_diff = FiniteDifferences.grad(central_fdm(5, 1), x -> tcholsolv_lower(x, B, i), x)[1] @test reverse_grad ≈ finite_diff #@test forward_grad ≈ finite_diff # - reverse_grad = Enzyme.gradient(Reverse, Const(x -> tcholsolv_upper(x, B, i)), x) - #forward_grad = reshape(collect(Enzyme.gradient(Forward, x -> tcholsolv_upper(x, B, i), x)), size(x)) + reverse_grad = Enzyme.gradient(Reverse, Const(x -> tcholsolv_upper(x, B, i)), x)[1] + #forward_grad = Enzyme.gradient(Forward, x -> tcholsolv_upper(x, B, i), x)[1] finite_diff = FiniteDifferences.grad(central_fdm(5, 1), x -> tcholsolv_upper(x, B, i), x)[1] @test reverse_grad ≈ finite_diff #@test forward_grad ≈ finite_diff @@ -554,7 +554,7 @@ end b = [1., 2.] dA = zero(A) Enzyme.autodiff(Reverse, h, Active, Duplicated(A, dA), Const(b)) - # dA_fwd = Enzyme.gradient(Forward, A->h(A, b), A) + # dA_fwd = Enzyme.gradient(Forward, A->h(A, b), A)[1] dA_fd = FiniteDifferences.grad(central_fdm(5, 1), A->h(A, b), A)[1] @test isapprox(dA, dA_fd) @@ -571,9 +571,9 @@ end @testset "Cholesky upper triangular v1" begin x = [1.0, -0.10541615131279458, 0.6219810761363638, 0.293343219811946, -0.10541615131279458, 1.0, -0.05258941747718969, 0.34629296878264443, 0.6219810761363638, -0.05258941747718969, 1.0, 0.4692436399208845, 0.293343219811946, 0.34629296878264443, 0.4692436399208845, 1.0] - @test collect(Enzyme.gradient(Forward, chol_upper, x)) ≈ [0.05270807565639728, 0.0, 0.0, 0.0, 0.9999999999999999, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + @test Enzyme.gradient(Forward, chol_upper, x)[1] ≈ [0.05270807565639728, 0.0, 0.0, 0.0, 0.9999999999999999, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] - @test Enzyme.gradient(Reverse, chol_upper, x) ≈ [0.05270807565639728, 0.0, 0.0, 0.0, 0.9999999999999999, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + @test Enzyme.gradient(Reverse, chol_upper, x)[1] ≈ [0.05270807565639728, 0.0, 0.0, 0.0, 0.9999999999999999, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] end @testset "Linear solve for triangular matrices" begin diff --git a/test/runtests.jl b/test/runtests.jl index b079c0f540..65ad4e3fd4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -488,9 +488,9 @@ end @testset "Deferred upgrade" begin function gradsin(x) - return gradient(Reverse, sin, x) + return gradient(Reverse, sin, x)[1] end - res = Enzyme.gradient(Reverse, gradsin, 3.1) + res = Enzyme.gradient(Reverse, gradsin, 3.1)[1] @test res ≈ -sin(3.1) end @@ -2794,43 +2794,43 @@ end @testset "Gradient & NamedTuples" begin xy = (x = [1.0, 2.0], y = [3.0, 4.0]) - grad = Enzyme.gradient(Reverse, z -> sum(z.x .* z.y), xy) + grad = Enzyme.gradient(Reverse, z -> sum(z.x .* z.y), xy)[1] @test grad == (x = [3.0, 4.0], y = [1.0, 2.0]) xp = (x = [1.0, 2.0], p = 3) # 3::Int is non-diff - grad = Enzyme.gradient(Reverse, z -> sum(z.x .^ z.p), xp) + grad = Enzyme.gradient(Reverse, z -> sum(z.x .^ z.p), xp)[1] @test grad.x == [3.0, 12.0] xp2 = (x = [1.0, 2.0], p = 3.0) # mixed activity - grad = Enzyme.gradient(Reverse, z -> sum(z.x .^ z.p), xp2) + grad = Enzyme.gradient(Reverse, z -> sum(z.x .^ z.p), xp2)[1] @test grad.x == [3.0, 12.0] @test grad.p ≈ 5.545177444479562 xy = (x = [1.0, 2.0], y = [3, 4]) # y is non-diff - grad = Enzyme.gradient(Reverse, z -> sum(z.x .* z.y), xy) + grad = Enzyme.gradient(Reverse, z -> sum(z.x .* z.y), xy)[1] @test grad.x == [3.0, 4.0] @test grad.y === xy.y # make_zero did not copy this - grad = Enzyme.gradient(Reverse, z -> (z.x * z.y), (x=5.0, y=6.0)) + grad = Enzyme.gradient(Reverse, z -> (z.x * z.y), (x=5.0, y=6.0))[1] @test grad == (x = 6.0, y = 5.0) - grad = Enzyme.gradient(Reverse, abs2, 7.0) + grad = Enzyme.gradient(Reverse, abs2, 7.0)[1] @test grad == 14.0 end @testset "Gradient & SparseArrays / StaticArrays" begin x = sparse([5.0, 0.0, 6.0]) - dx = Enzyme.gradient(Reverse, sum, x) + dx = Enzyme.gradient(Reverse, sum, x)[1] @test dx isa SparseVector @test dx ≈ [1, 0, 1] x = sparse([5.0 0.0 6.0]) - dx = Enzyme.gradient(Reverse, sum, x) + dx = Enzyme.gradient(Reverse, sum, x)[1] @test dx isa SparseMatrixCSC @test dx ≈ [1 0 1] x = @SArray [5.0 0.0 6.0] - dx = Enzyme.gradient(Reverse, prod, x) + dx = Enzyme.gradient(Reverse, prod, x)[1] @test dx isa SArray @test dx ≈ [0 30 0] @@ -2851,7 +2851,7 @@ end @test y[2] == [0.0, 0.0, 1.0] x = @SArray [5.0 0.0 6.0] - dx = Enzyme.gradient(Forward, prod, x) + dx = Enzyme.gradient(Forward, prod, x)[1] @test dx[1] ≈ 0 @test dx[2] ≈ 30 @test dx[3] ≈ 0 @@ -2906,264 +2906,266 @@ mkarray(sz, args...) = reshape(vcat(args...), sz) scalar = 3.0 # ∂ scalar / ∂ scalar - @test Enzyme.gradient(Enzyme.Forward, x -> x^2, scalar) ≈ 6.0 - @test Enzyme.gradient(Enzyme.Reverse, x -> x^2, scalar) ≈ 6.0 - @test Enzyme.jacobian(Enzyme.Forward, x -> x^2, scalar) ≈ 6.0 - @test Enzyme.jacobian(Enzyme.Reverse, x -> x^2, scalar) ≈ 6.0 - @test Enzyme.gradient(Enzyme.Forward, x -> 2*x, scalar) ≈ 2.0 - @test Enzyme.gradient(Enzyme.Reverse, x -> 2*x, scalar) ≈ 2.0 - @test Enzyme.jacobian(Enzyme.Forward, x -> 2*x, scalar) ≈ 2.0 - @test Enzyme.jacobian(Enzyme.Reverse, x -> 2*x, scalar) ≈ 2.0 + @test Enzyme.gradient(Enzyme.Forward, x -> x^2, scalar)[1] ≈ 6.0 + @test Enzyme.gradient(Enzyme.Reverse, x -> x^2, scalar)[1] ≈ 6.0 + @test Enzyme.jacobian(Enzyme.Forward, x -> x^2, scalar)[1] ≈ 6.0 + @test Enzyme.jacobian(Enzyme.Reverse, x -> x^2, scalar)[1] ≈ 6.0 + @test Enzyme.gradient(Enzyme.Forward, x -> 2*x, scalar)[1] ≈ 2.0 + @test Enzyme.gradient(Enzyme.Reverse, x -> 2*x, scalar)[1] ≈ 2.0 + @test Enzyme.jacobian(Enzyme.Forward, x -> 2*x, scalar)[1] ≈ 2.0 + @test Enzyme.jacobian(Enzyme.Reverse, x -> 2*x, scalar)[1] ≈ 2.0 # ∂ vector / ∂ scalar - @test Enzyme.gradient(Enzyme.Forward, x -> [2*x, x^2], scalar) ≈ [2.0, 6.0] - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [2*x, x^2], scalar) ≈ [2.0, 6.0] + @test Enzyme.gradient(Enzyme.Forward, x -> [2*x, x^2], scalar)[1] ≈ [2.0, 6.0] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [2*x, x^2], scalar)[1] ≈ [2.0, 6.0] - @test Enzyme.jacobian(Enzyme.Forward, x -> [2*x, x^2], scalar) ≈ [2.0, 6.0] - @test Enzyme.jacobian(Enzyme.Reverse, x -> [2*x, x^2], scalar) ≈ [2.0, 6.0] + @test Enzyme.jacobian(Enzyme.Forward, x -> [2*x, x^2], scalar)[1] ≈ [2.0, 6.0] + @test Enzyme.jacobian(Enzyme.Reverse, x -> [2*x, x^2], scalar)[1] ≈ [2.0, 6.0] # ∂ tuple / ∂ scalar - @test Enzyme.gradient(Enzyme.Forward, x -> (2*x, x^2), scalar) ≃ (2.0, 6.0) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (2*x, x^2), scalar) ≈ [2.0, 6.0] + @test Enzyme.gradient(Enzyme.Forward, x -> (2*x, x^2), scalar)[1] ≃ (2.0, 6.0) + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (2*x, x^2), scalar)[1] ≈ [2.0, 6.0] - @test Enzyme.jacobian(Enzyme.Forward, x -> (2*x, x^2), scalar) ≃ (2.0, 6.0) - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (2*x, x^2), scalar) ≃ (2.0, 6.0) + @test Enzyme.jacobian(Enzyme.Forward, x -> (2*x, x^2), scalar)[1] ≃ (2.0, 6.0) + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (2*x, x^2), scalar)[1] ≃ (2.0, 6.0) mkarray1 = x -> mkarray((2,2),2*x,sin(x),x^2,exp(x)) # ∂ matrix / ∂ scalar - @test Enzyme.gradient(Enzyme.Forward, mkarray1, scalar) ≈ [2.0 6.0; cos(scalar) exp(scalar)] - @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray1, scalar) ≈ [2.0 6.0; cos(scalar) exp(scalar)] + @test Enzyme.gradient(Enzyme.Forward, mkarray1, scalar)[1] ≈ [2.0 6.0; cos(scalar) exp(scalar)] + @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray1, scalar)[1] ≈ [2.0 6.0; cos(scalar) exp(scalar)] - @test Enzyme.jacobian(Enzyme.Forward, mkarray1, scalar) ≈ [2.0 6.0; cos(scalar) exp(scalar)] - @test Enzyme.jacobian(Enzyme.Reverse, mkarray1, scalar) ≈ [2.0 6.0; cos(scalar) exp(scalar)] + @test Enzyme.jacobian(Enzyme.Forward, mkarray1, scalar)[1] ≈ [2.0 6.0; cos(scalar) exp(scalar)] + @test Enzyme.jacobian(Enzyme.Reverse, mkarray1, scalar)[1] ≈ [2.0 6.0; cos(scalar) exp(scalar)] # ∂ struct / ∂ scalar - @test Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x, x^2, x^3), scalar) == OutStruct(1.0,2*scalar,3*scalar^2) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> InpStruct(x, x^2, x^3), scalar) == (OutStruct(1.0,2.0,3.0),) - @test Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x, x^2, x^3), scalar) == OutStruct(1.0,2*scalar,3*scalar^2) - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> InpStruct(x, x^2, x^3), scalar) == (OutStruct(1.0,2.0,3.0),) + @test Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x, x^2, x^3), scalar)[1] == OutStruct(1.0,2*scalar,3*scalar^2) + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> InpStruct(x, x^2, x^3), scalar)[1] == (OutStruct(1.0,2.0,3.0),) + @test Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x, x^2, x^3), scalar)[1] == OutStruct(1.0,2*scalar,3*scalar^2) + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> InpStruct(x, x^2, x^3), scalar)[1] == (OutStruct(1.0,2.0,3.0),) vector = [2.7, 3.1] # ∂ scalar / ∂ vector - @test Enzyme.gradient(Enzyme.Forward, x -> x[1] * x[2], vector) ≃ (vector[2],vector[1]) - @test Enzyme.gradient(Enzyme.Reverse, x -> x[1] * x[2], vector) ≈ [vector[2], vector[1]] - @test Enzyme.jacobian(Enzyme.Forward, x -> x[1] * x[2], vector) ≈ [vector[2], vector[1]] - @test Enzyme.jacobian(Enzyme.Reverse, x -> x[1] * x[2], vector) ≈ [vector[2], vector[1]] + @test Enzyme.gradient(Enzyme.Forward, x -> x[1] * x[2], vector)[1] ≈ [vector[2],vector[1]] + @test Enzyme.gradient(Enzyme.Reverse, x -> x[1] * x[2], vector)[1] ≈ [vector[2], vector[1]] + @test Enzyme.jacobian(Enzyme.Forward, x -> x[1] * x[2], vector)[1] ≈ [vector[2], vector[1]] + @test Enzyme.jacobian(Enzyme.Reverse, x -> x[1] * x[2], vector)[1] ≈ [vector[2], vector[1]] # ∂ vector / ∂ vector - @test Enzyme.gradient(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector) ≃ - ([vector[2], -sin(vector[1])], [vector[1], 1.0]) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector) ≈ ([vector[2], -sin(vector[1])], [vector[1], 1.0]) - @test Enzyme.jacobian(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector) ≈ + @test Enzyme.gradient(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector)[1] ≈ + [vector[2] vector[1]; -sin(vector[1]) 1.0] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector)[1] ≈ + [vector[2] vector[1]; -sin(vector[1]) 1.0] + @test Enzyme.jacobian(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector)[1] ≈ [vector[2] vector[1]; -sin(vector[1]) 1.0] - @test Enzyme.jacobian(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector) ≈ + @test Enzyme.jacobian(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], vector)[1] ≈ [vector[2] vector[1]; -sin(vector[1]) 1.0] # ∂ tuple / ∂ vector - @test Enzyme.gradient(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector) ≃ - ((vector[2], -sin(vector[1])), (vector[1], 1.0)) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector) ≈ + @test Enzyme.gradient(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector)[1] ≃ + [(vector[2], -sin(vector[1])), (vector[1], 1.0)] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector)[1] ≈ ([vector[2], -sin(vector[1])], [vector[1], 1.0]) - @test Enzyme.jacobian(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector) ≃ + @test Enzyme.jacobian(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector)[1] ≃ [(vector[2], -sin(vector[1])), (vector[1], 1.0)] - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector) + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector)[1] mkarray2 = x -> mkarray((2,2), x[1]*x[2], exp(x[2]), cos(x[1])+x[2], x[1]) # ∂ matrix / ∂ vector - @test Enzyme.gradient(Enzyme.Forward, mkarray2, vector) ≃ - ([vector[2] -sin(vector[1]); 0.0 1.0], [vector[1] 1.0; exp(vector[2]) 0.0]) - @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray2, vector) - @test Enzyme.jacobian(Enzyme.Forward, mkarray2, vector) ≈ + @test Enzyme.gradient(Enzyme.Forward, mkarray2, vector)[1] ≈ + mkarray((2,2,2), vector[2], 0.0, -sin(vector[1]), 1.0, vector[1], exp(vector[2]), 1.0, 0.0) + @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray2, vector)[1] + @test Enzyme.jacobian(Enzyme.Forward, mkarray2, vector)[1] ≈ mkarray((2,2,2), vector[2], 0.0, -sin(vector[1]), 1.0, vector[1], exp(vector[2]), 1.0, 0.0) - @test Enzyme.jacobian(Enzyme.Reverse, mkarray2, vector) ≈ + @test Enzyme.jacobian(Enzyme.Reverse, mkarray2, vector)[1] ≈ mkarray((2,2,2), vector[2], 0.0, -sin(vector[1]), 1.0, vector[1], exp(vector[2]), 1.0, 0.0) # ∂ struct / ∂ vector - @test Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), vector) ≃ - (OutStruct(vector[2], -sin(vector[1]), 0.0), OutStruct(vector[1], 1.0, exp(vector[2]))) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector) ≈ ([vector[2], -sin(vector[1])], [vector[1], 1.0]) + @test Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), vector)[1] ≃ + [OutStruct(vector[2], -sin(vector[1]), 0.0), OutStruct(vector[1], 1.0, exp(vector[2]))] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector)[1] ≈ ([vector[2], -sin(vector[1])], [vector[1], 1.0]) - @test Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), vector) ≃ + @test Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), vector)[1] ≃ [OutStruct(vector[2], -sin(vector[1]), 0.0), OutStruct(vector[1], 1.0, exp(vector[2]))] - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector) ≈ ([vector[2], -sin(vector[1])], [vector[1], 1.0]) + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), vector)[1] ≈ ([vector[2], -sin(vector[1])], [vector[1], 1.0]) tuplev = (2.7, 3.1) # ∂ scalar / ∂ tuple - @test Enzyme.gradient(Enzyme.Forward, x -> x[1] * x[2], tuplev) ≃ (tuplev[2],tuplev[1]) - @test Enzyme.gradient(Enzyme.Reverse, x -> x[1] * x[2], tuplev) ≃ (tuplev[2],tuplev[1]) - @test Enzyme.jacobian(Enzyme.Forward, x -> x[1] * x[2], tuplev) ≃ (tuplev[2],tuplev[1]) - @test Enzyme.jacobian(Enzyme.Reverse, x -> x[1] * x[2], tuplev) ≃ (tuplev[2],tuplev[1]) + @test Enzyme.gradient(Enzyme.Forward, x -> x[1] * x[2], tuplev)[1] ≃ (tuplev[2],tuplev[1]) + @test Enzyme.gradient(Enzyme.Reverse, x -> x[1] * x[2], tuplev)[1] ≃ (tuplev[2],tuplev[1]) + @test Enzyme.jacobian(Enzyme.Forward, x -> x[1] * x[2], tuplev)[1] ≃ (tuplev[2],tuplev[1]) + @test Enzyme.jacobian(Enzyme.Reverse, x -> x[1] * x[2], tuplev)[1] ≃ (tuplev[2],tuplev[1]) # ∂ vector / ∂ tuple - @test Enzyme.gradient(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev) ≃ + @test Enzyme.gradient(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev)[1] ≃ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev) ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) - @test_broken Enzyme.jacobian(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev) ≈ + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev)[1] ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) + @test_broken Enzyme.jacobian(Enzyme.Forward, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev)[1] ≈ [tuplev[2] tuplev[1]; -sin(tuplev[1]) 1.0] - @test Enzyme.jacobian(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev) ≃ + @test Enzyme.jacobian(Enzyme.Reverse, x -> [x[1] * x[2], cos(x[1]) + x[2]], tuplev)[1] ≃ [(tuplev[2], tuplev[1]), (-sin(tuplev[1]), 1.0)] # ∂ tuple / ∂ tuple - @test Enzyme.gradient(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev) ≃ + @test Enzyme.gradient(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev)[1] ≃ ((vector[2], -sin(vector[1])), (vector[1], 1.0)) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev) ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) - @test Enzyme.jacobian(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev) ≃ + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev)[1] ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) + @test Enzyme.jacobian(Enzyme.Forward, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev)[1] ≃ ((tuplev[2], -sin(tuplev[1])), (tuplev[1], 1.0)) - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev) ≈ + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev)[1] ≈ [tuplev[2] tuplev[1]; -sin(tuplev[1]) 1.0] # ∂ matrix / ∂ tuple - @test Enzyme.gradient(Enzyme.Forward, mkarray2, tuplev) ≃ + @test Enzyme.gradient(Enzyme.Forward, mkarray2, tuplev)[1] ≃ ([tuplev[2] -sin(tuplev[1]); 0.0 1.0], [tuplev[1] 1.0; exp(tuplev[2]) 0.0]) - @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray2, tuplev) - @test_broken Enzyme.jacobian(Enzyme.Forward, mkarray2, tuplev) ≈ + @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray2, tuplev)[1] + @test_broken Enzyme.jacobian(Enzyme.Forward, mkarray2, tuplev)[1] ≈ [tuplev[2] -sin(tuplev[1]); 0.0 1.0;;; tuplev[1] 1.0; exp(tuplev[2]) 0.0] - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> mkarray2, tuplev) ≈ + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> mkarray2, tuplev)[1] ≈ [tuplev[2] -sin(tuplev[1]); 0.0 1.0;;; tuplev[1] 1.0; exp(tuplev[2]) 0.0] # ∂ struct / ∂ tuple - @test Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), tuplev) ≃ + @test Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), tuplev)[1] ≃ (OutStruct(tuplev[2], -sin(tuplev[1]), 0.0), OutStruct(tuplev[1], 1.0, exp(tuplev[2]))) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev) ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev)[1] ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) - @test_broken Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), tuplev) ≃ + @test_broken Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x[1] * x[2], cos(x[1]) + x[2], exp(x[2])), tuplev)[1] ≃ [OutStruct(tuplev[2], -sin(tuplev[1]), 0.0), OutStruct(tuplev[1], 1.0, exp(tuplev[2]))] - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev) ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x[1] * x[2], cos(x[1]) + x[2]), tuplev)[1] ≈ ([tuplev[2], -sin(tuplev[1])], [tuplev[1], 1.0]) matrix = [2.7 3.1; 4.7 5.6] # ∂ scalar / ∂ matrix - @test Enzyme.gradient(Enzyme.Forward, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix) ≃ - (matrix[1,2], matrix[2,2], matrix[1,1], matrix[2,1]) - @test Enzyme.gradient(Enzyme.Reverse, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix) ≈ [matrix[1,2] matrix[1,1]; matrix[2,2] matrix[2,1]] - @test Enzyme.jacobian(Enzyme.Forward, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix) ≈ [matrix[1,2] matrix[1,1]; matrix[2,2] matrix[2,1]] - @test Enzyme.jacobian(Enzyme.Reverse, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix) ≈ [matrix[1,2] matrix[1,1]; matrix[2,2] matrix[2,1]] + @test Enzyme.gradient(Enzyme.Forward, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix)[1] ≈ [matrix[1,2] matrix[1,1]; matrix[2,2] matrix[2,1]] + @test Enzyme.gradient(Enzyme.Reverse, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix)[1] ≈ [matrix[1,2] matrix[1,1]; matrix[2,2] matrix[2,1]] + @test Enzyme.jacobian(Enzyme.Forward, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix)[1] ≈ [matrix[1,2] matrix[1,1]; matrix[2,2] matrix[2,1]] + @test Enzyme.jacobian(Enzyme.Reverse, x->x[1,1]*x[1,2]+x[2,1]*x[2,2], matrix)[1] ≈ [matrix[1,2] matrix[1,1]; matrix[2,2] matrix[2,1]] # ∂ vector / ∂ matrix - @test Enzyme.gradient(Enzyme.Forward, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix) ≃ - ([matrix[1,2], 0.0], [0.0, matrix[2,2]], [matrix[1,1], 0.0], [0.0, matrix[2,1]]) - @test_broken Enzyme.gradient(Enzyme.Reverse, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix) + @test Enzyme.gradient(Enzyme.Forward, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix)[1] ≈ + mkarray((2,2,2), matrix[1,2], 0.0, 0.0, matrix[2,2], matrix[1,1], 0.0, 0.0, matrix[2,1]) + @test_broken Enzyme.gradient(Enzyme.Reverse, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix)[1] # again we can't use array construction syntax because of 1.6 - @test Enzyme.jacobian(Enzyme.Forward, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix) ≈ + @test Enzyme.jacobian(Enzyme.Forward, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix)[1] ≈ mkarray((2,2,2), matrix[1,2], 0.0, 0.0, matrix[2,2], matrix[1,1], 0.0, 0.0, matrix[2,1]) - @test Enzyme.jacobian(Enzyme.Reverse, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix) ≈ + @test Enzyme.jacobian(Enzyme.Reverse, x->[x[1,1]*x[1,2],x[2,1]*x[2,2]], matrix)[1] ≈ mkarray((2,2,2), matrix[1,2], 0.0, 0.0, matrix[2,2], matrix[1,1], 0.0, 0.0, matrix[2,1]) # ∂ tuple / ∂ matrix - @test Enzyme.gradient(Enzyme.Forward, x->(x[1,1]*x[1,2],x[2,1]*x[2,2]), matrix) ≃ ((matrix[1,2], 0.0), (0.0, matrix[2,2]), (matrix[1,1], 0.0), (0.0, matrix[2,1])) + @test Enzyme.gradient(Enzyme.Forward, x->(x[1,1]*x[1,2],x[2,1]*x[2,2]), matrix)[1] ≃ + [(matrix[1,2],0.0) (matrix[1,1],0.0); (0.0,matrix[2,2]) (0.0,matrix[2,1])] @test_broken Enzyme.gradient(Enzyme.Reverse, x->(x[1,1]*x[1,2],x[2,1]*x[2,2]), matrix) - @test Enzyme.jacobian(Enzyme.Forward, x->(x[1,1]*x[1,2],x[2,1]*x[2,2]), matrix) ≃ + @test Enzyme.jacobian(Enzyme.Forward, x->(x[1,1]*x[1,2],x[2,1]*x[2,2]), matrix)[1] ≃ [(matrix[1,2],0.0) (matrix[1,1],0.0); (0.0,matrix[2,2]) (0.0,matrix[2,1])] - @test_broken Enzyme.jacobian(Enzyme.Reverse, x->(x[1,1]*x[1,2],x[2,1]*x[2,2]), matrix) + @test_broken Enzyme.jacobian(Enzyme.Reverse, x->(x[1,1]*x[1,2],x[2,1]*x[2,2]), matrix)[1] mkarray3 = x -> mkarray((2,2), x[1,1]*x[1,2], exp(x[1,1])+x[2,2], x[2,1]*x[2,2], sin(x[1,2])+x[2,1]) # ∂ matrix / ∂ matrix - @test Enzyme.gradient(Enzyme.Forward, mkarray3, matrix) ≃ - ([matrix[1,2] 0.0; exp(matrix[1,1]) 0.0], [0.0 matrix[2,2]; 0.0 1.0], [matrix[1,1] 0.0; 0.0 cos(matrix[1,2])], [0.0 matrix[2,1]; 1.0 0.0]) - @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray3, matrix) + @test Enzyme.gradient(Enzyme.Forward, mkarray3, matrix)[1] ≈ + mkarray((2,2,2,2), matrix[1,2],exp(matrix[1,1]),0.0,0.0,0.0,0.0,matrix[2,2],1.0, + matrix[1,1],0.0,0.0,cos(matrix[1,2]),0.0,1.0,matrix[2,1],0.0) + @test_broken Enzyme.gradient(Enzyme.Reverse, mkarray3, matrix)[1] # array construction syntax broken on 1.6 - @test Enzyme.jacobian(Enzyme.Forward, mkarray3, matrix) ≈ + @test Enzyme.jacobian(Enzyme.Forward, mkarray3, matrix)[1] ≈ mkarray((2,2,2,2), matrix[1,2],exp(matrix[1,1]),0.0,0.0,0.0,0.0,matrix[2,2],1.0, matrix[1,1],0.0,0.0,cos(matrix[1,2]),0.0,1.0,matrix[2,1],0.0) - @test Enzyme.jacobian(Enzyme.Reverse, mkarray3, matrix) ≈ + @test Enzyme.jacobian(Enzyme.Reverse, mkarray3, matrix)[1] ≈ mkarray((2,2,2,2), matrix[1,2],exp(matrix[1,1]),0.0,0.0,0.0,0.0,matrix[2,2],1.0, matrix[1,1],0.0,0.0,cos(matrix[1,2]),0.0,1.0,matrix[2,1],0.0) # ∂ tuple / ∂ matrix - @test Enzyme.gradient(Enzyme.Forward, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix) ≃ - (OutStruct(matrix[1,2], 0.0, exp(matrix[1,1])), OutStruct(0.0, matrix[2,2], 0.0), OutStruct(matrix[1,1], 0.0, 0.0), OutStruct(0.0, matrix[2,1], 1.0)) - @test_broken Enzyme.gradient(Enzyme.Reverse, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix) - @test Enzyme.jacobian(Enzyme.Forward, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix) ≃ + @test Enzyme.gradient(Enzyme.Forward, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix)[1] ≃ + [OutStruct(matrix[1,2],0.0, exp(matrix[1,1])) OutStruct(matrix[1,1],0.0,0.0); OutStruct(0.0,matrix[2,2],0.0) OutStruct(0.0,matrix[2,1], 1.0)] + @test_broken Enzyme.gradient(Enzyme.Reverse, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix)[1] + @test Enzyme.jacobian(Enzyme.Forward, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix)[1] ≃ [OutStruct(matrix[1,2],0.0, exp(matrix[1,1])) OutStruct(matrix[1,1],0.0,0.0); OutStruct(0.0,matrix[2,2],0.0) OutStruct(0.0,matrix[2,1], 1.0)] - @test_broken Enzyme.jacobian(Enzyme.Reverse, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix) + @test_broken Enzyme.jacobian(Enzyme.Reverse, x->OutStruct(x[1,1]*x[1,2],x[2,1]*x[2,2], exp(x[1,1])+x[2,2]), matrix)[1] istruct = InpStruct(2.7, 3.1, 4.7) # ∂ scalar / ∂ struct - @test_broken Enzyme.gradient(Enzyme.Forward, x -> x.i1 * x.i2 + x.i3, istruct) - @test Enzyme.gradient(Enzyme.Reverse, x -> x.i1 * x.i2 + x.i3, istruct) ≃ InpStruct(istruct.i2, istruct.i1, 1.0) - @test_broken Enzyme.jacobian(Enzyme.Forward, x -> x.i1 * x.i2 + x.i3, istruct) - @test Enzyme.jacobian(Enzyme.Reverse, x -> x.i1 * x.i2 + x.i3, istruct) ≃ InpStruct(istruct.i2, istruct.i1, 1.0) + @test_broken Enzyme.gradient(Enzyme.Forward, x -> x.i1 * x.i2 + x.i3, istruct)[1] + @test Enzyme.gradient(Enzyme.Reverse, x -> x.i1 * x.i2 + x.i3, istruct)[1] ≃ InpStruct(istruct.i2, istruct.i1, 1.0) + @test_broken Enzyme.jacobian(Enzyme.Forward, x -> x.i1 * x.i2 + x.i3, istruct)[1] + @test Enzyme.jacobian(Enzyme.Reverse, x -> x.i1 * x.i2 + x.i3, istruct)[1] ≃ InpStruct(istruct.i2, istruct.i1, 1.0) # ∂ vector / ∂ struct - @test_broken Enzyme.gradient(Enzyme.Forward, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct) - @test_broken Enzyme.jacobian(Enzyme.Forward, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct) - @test Enzyme.jacobian(Enzyme.Reverse, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct) ≃ [InpStruct(istruct.i2, istruct.i1, 0.0), InpStruct(1.0, 0.0, -sin(istruct.i3))] + @test_broken Enzyme.gradient(Enzyme.Forward, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct)[1] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct)[1] + @test_broken Enzyme.jacobian(Enzyme.Forward, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct)[1] + @test Enzyme.jacobian(Enzyme.Reverse, x -> [x.i1 * x.i2, cos(x.i3) + x.i1], istruct)[1] ≃ [InpStruct(istruct.i2, istruct.i1, 0.0), InpStruct(1.0, 0.0, -sin(istruct.i3))] # ∂ tuple / ∂ struct - @test_broken Enzyme.gradient(Enzyme.Forward, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct) - @test_broken Enzyme.jacobian(Enzyme.Forward, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct) - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct) + @test_broken Enzyme.gradient(Enzyme.Forward, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct)[1] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct)[1] + @test_broken Enzyme.jacobian(Enzyme.Forward, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct)[1] + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> (x.i1 * x.i2, cos(x.i3) + x.i1), istruct)[1] mkarray4 = x -> mkarray((2,2), x.i1*x.i2, exp(x.i2), cos(x.i3)+x.i1, x.i1) # ∂ matrix / ∂ struct - @test_broken Enzyme.gradient(Enzyme.Forward, x -> [x.i1 * x.i2 cos(x.i3) + x.i1; exp(x.i2) x.i1], istruct) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x.i1 * x.i2 cos(x.i3) + x.i1; exp(x.i2) x.i1], istruct) - @test_broken Enzyme.jacobian(Enzyme.Forward, x -> [x.i1 * x.i2 cos(x.i3) + x.i1; exp(x.i2) x.i1], istruct) - @test Enzyme.jacobian(Enzyme.Reverse, mkarray4, istruct) ≃ + @test_broken Enzyme.gradient(Enzyme.Forward, x -> [x.i1 * x.i2 cos(x.i3) + x.i1; exp(x.i2) x.i1], istruct)[1] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> [x.i1 * x.i2 cos(x.i3) + x.i1; exp(x.i2) x.i1], istruct)[1] + @test_broken Enzyme.jacobian(Enzyme.Forward, x -> [x.i1 * x.i2 cos(x.i3) + x.i1; exp(x.i2) x.i1], istruct)[1] + @test Enzyme.jacobian(Enzyme.Reverse, mkarray4, istruct)[1] ≃ [InpStruct(istruct.i2, istruct.i1, 0.0) InpStruct(1.0, 0.0, -sin(istruct.i3)); InpStruct(0.0, exp(istruct.i2), 0.0) InpStruct(1.0, 0.0, 0.0)] # ∂ struct / ∂ struct - @test_broken Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct) - @test_broken Enzyme.gradient(Enzyme.Reverse, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct) - @test_broken Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct) - @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct) + @test_broken Enzyme.gradient(Enzyme.Forward, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct)[1] + @test_broken Enzyme.gradient(Enzyme.Reverse, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct)[1] + @test_broken Enzyme.jacobian(Enzyme.Forward, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct)[1] + @test_broken Enzyme.jacobian(Enzyme.Reverse, x -> OutStruct(x.i1 * x.i2, cos(x.i3) + x.i1, exp(x.i2)), istruct)[1] end @testset "Simple Jacobian" begin - @test Enzyme.jacobian(Enzyme.Forward, x->2*x, 3.0) ≈ 2.0 - @test Enzyme.jacobian(Enzyme.Forward, x->[x, 2*x], 3.0) ≈ [1.0, 2.0] - @test Enzyme.jacobian(Enzyme.Forward, x->sum(abs2, x), [2.0, 3.0]) ≈ [4.0, 6.0] + @test Enzyme.jacobian(Enzyme.Forward, x->2*x, 3.0)[1] ≈ 2.0 + @test Enzyme.jacobian(Enzyme.Forward, x->[x, 2*x], 3.0)[1] ≈ [1.0, 2.0] + @test Enzyme.jacobian(Enzyme.Forward, x->sum(abs2, x), [2.0, 3.0])[1] ≈ [4.0, 6.0] - @test Enzyme.jacobian(Enzyme.Forward, x->2*x, 3.0, Val(1)) ≈ 2.0 - @test Enzyme.jacobian(Enzyme.Forward, x->[x, 2*x], 3.0, Val(1)) ≈ [1.0, 2.0] - @test Enzyme.jacobian(Enzyme.Forward, x->sum(abs2, x), [2.0, 3.0], Val(1)) ≈ [4.0, 6.0] + @test Enzyme.jacobian(Enzyme.Forward, x->2*x, 3.0, chunk=Val(1))[1] ≈ 2.0 + @test Enzyme.jacobian(Enzyme.Forward, x->[x, 2*x], 3.0, chunk=Val(1))[1] ≈ [1.0, 2.0] + @test Enzyme.jacobian(Enzyme.Forward, x->sum(abs2, x), [2.0, 3.0], chunk=Val(1))[1] ≈ [4.0, 6.0] - @test Enzyme.jacobian(Enzyme.Forward, x->2*x, 3.0, Val(2)) ≈ 2.0 - @test Enzyme.jacobian(Enzyme.Forward, x->[x, 2*x], 3.0, Val(2)) ≈ [1.0, 2.0] - @test Enzyme.jacobian(Enzyme.Forward, x->sum(abs2, x), [2.0, 3.0], Val(2)) ≈ [4.0, 6.0] + @test Enzyme.jacobian(Enzyme.Forward, x->2*x, 3.0, chunk=Val(2))[1] ≈ 2.0 + @test Enzyme.jacobian(Enzyme.Forward, x->[x, 2*x], 3.0, chunk=Val(2))[1] ≈ [1.0, 2.0] + @test Enzyme.jacobian(Enzyme.Forward, x->sum(abs2, x), [2.0, 3.0], chunk=Val(2))[1] ≈ [4.0, 6.0] - @test Enzyme.jacobian(Enzyme.Reverse, x->[x, 2*x], 3.0, Val(2)) ≈ [1.0, 2.0] - @test Enzyme.jacobian(Enzyme.Reverse, x->[x, 2*x], 3.0, Val(2), Val(1)) ≈ [1.0, 2.0] - @test Enzyme.jacobian(Enzyme.Reverse, x->[x, 2*x], 3.0, Val(2), Val(2)) ≈ [1.0, 2.0] + @test Enzyme.jacobian(Enzyme.Reverse, x->[x, 2*x], 3.0, n_outs=Val((2,)))[1] ≈ [1.0, 2.0] + @test Enzyme.jacobian(Enzyme.Reverse, x->[x, 2*x], 3.0, n_outs=Val((2,)), chunk=Val(1))[1] ≈ [1.0, 2.0] + @test Enzyme.jacobian(Enzyme.Reverse, x->[x, 2*x], 3.0, n_outs=Val((2,)), chunk=Val(2))[1] ≈ [1.0, 2.0] x = float.(reshape(1:6, 2, 3)) fillabs2(x) = [sum(abs2, x), 10*sum(abs2, x), 100*sum(abs2, x), 1000*sum(abs2, x)] - jac = Enzyme.jacobian(Enzyme.Forward, fillabs2, x) + jac = Enzyme.jacobian(Enzyme.Forward, fillabs2, x)[1] @test jac[1, :, :] ≈ [2.0 6.0 10.0; 4.0 8.0 12.0] @test jac[2, :, :] ≈ [20.0 60.0 100.0; 40.0 80.0 120.0] @test jac[3, :, :] ≈ [200.0 600.0 1000.0; 400.0 800.0 1200.0] @test jac[4, :, :] ≈ [2000.0 6000.0 10000.0; 4000.0 8000.0 12000.0] - jac = Enzyme.jacobian(Enzyme.Forward, fillabs2, x, Val(1)) + jac = Enzyme.jacobian(Enzyme.Forward, fillabs2, x, chunk=Val(1))[1] @test jac[1, :, :] ≈ [2.0 6.0 10.0; 4.0 8.0 12.0] @test jac[2, :, :] ≈ [20.0 60.0 100.0; 40.0 80.0 120.0] @test jac[3, :, :] ≈ [200.0 600.0 1000.0; 400.0 800.0 1200.0] @test jac[4, :, :] ≈ [2000.0 6000.0 10000.0; 4000.0 8000.0 12000.0] - jac = Enzyme.jacobian(Enzyme.Forward, fillabs2, x, Val(2)) + jac = Enzyme.jacobian(Enzyme.Forward, fillabs2, x, chunk=Val(2))[1] @test jac[1, :, :] ≈ [2.0 6.0 10.0; 4.0 8.0 12.0] @test jac[2, :, :] ≈ [20.0 60.0 100.0; 40.0 80.0 120.0] @@ -3171,14 +3173,14 @@ end @test jac[4, :, :] ≈ [2000.0 6000.0 10000.0; 4000.0 8000.0 12000.0] - jac = Enzyme.jacobian(Enzyme.Reverse, fillabs2, x, Val(4), Val(1)) + jac = Enzyme.jacobian(Enzyme.Reverse, fillabs2, x, n_outs=Val((4,)), chunk=Val(1))[1] @test jac[1, :, :] ≈ [2.0 6.0 10.0; 4.0 8.0 12.0] @test jac[2, :, :] ≈ [20.0 60.0 100.0; 40.0 80.0 120.0] @test jac[3, :, :] ≈ [200.0 600.0 1000.0; 400.0 800.0 1200.0] @test jac[4, :, :] ≈ [2000.0 6000.0 10000.0; 4000.0 8000.0 12000.0] - jac = Enzyme.jacobian(Enzyme.Reverse, fillabs2, x, Val(4), Val(2)) + jac = Enzyme.jacobian(Enzyme.Reverse, fillabs2, x, n_outs=Val((4,)), chunk=Val(2))[1] @test jac[1, :, :] ≈ [2.0 6.0 10.0; 4.0 8.0 12.0] @test jac[2, :, :] ≈ [20.0 60.0 100.0; 40.0 80.0 120.0] @@ -3189,14 +3191,14 @@ end x2 = InpStruct(1.0, 2.0, 3.0) - jac = Enzyme.jacobian(Enzyme.Reverse, fillinpabs2, x2, Val(4), Val(1)) + jac = Enzyme.jacobian(Enzyme.Reverse, fillinpabs2, x2, n_outs=Val((4,)), chunk=Val(1))[1] @test jac[1] == InpStruct(2.0, 4.0, 6.0) @test jac[2] == InpStruct(20.0, 40.0, 60.0) @test jac[3] == InpStruct(200.0, 400.0, 600.0) @test jac[4] == InpStruct(2000.0, 4000.0, 6000.0) - jac = Enzyme.jacobian(Enzyme.Reverse, fillinpabs2, x2, Val(4), Val(2)) + jac = Enzyme.jacobian(Enzyme.Reverse, fillinpabs2, x2, n_outs=Val((4,)), chunk=Val(2))[1] @test jac[1] == InpStruct(2.0, 4.0, 6.0) @test jac[2] == InpStruct(20.0, 40.0, 60.0) @@ -3205,7 +3207,7 @@ end filloutabs2(x) = OutStruct(sum(abs2, x), 10*sum(abs2, x), 100*sum(abs2, x)) - jac = Enzyme.jacobian(Enzyme.Forward, filloutabs2, x) + jac = Enzyme.jacobian(Enzyme.Forward, filloutabs2, x)[1] @test jac[1, 1] == OutStruct(2.0, 20.0, 200.0) @test jac[2, 1] == OutStruct(4.0, 40.0, 400.0) @@ -3216,7 +3218,7 @@ end @test jac[1, 3] == OutStruct(10.0, 100.0, 1000.0) @test jac[2, 3] == OutStruct(12.0, 120.0, 1200.0) - jac = Enzyme.jacobian(Enzyme.Forward, filloutabs2, x, Val(1)) + jac = Enzyme.jacobian(Enzyme.Forward, filloutabs2, x, chunk=Val(1))[1] @test jac[1, 1] == OutStruct(2.0, 20.0, 200.0) @test jac[2, 1] == OutStruct(4.0, 40.0, 400.0) @@ -3227,7 +3229,7 @@ end @test jac[1, 3] == OutStruct(10.0, 100.0, 1000.0) @test jac[2, 3] == OutStruct(12.0, 120.0, 1200.0) - jac = Enzyme.jacobian(Enzyme.Forward, filloutabs2, x, Val(2)) + jac = Enzyme.jacobian(Enzyme.Forward, filloutabs2, x, chunk=Val(2))[1] @test jac[1, 1] == OutStruct(2.0, 20.0, 200.0) @test jac[2, 1] == OutStruct(4.0, 40.0, 400.0) @@ -3245,27 +3247,27 @@ end [v[2], v[1]*v[1], v[1]*v[1]*v[1]] end - jac = Enzyme.jacobian(Reverse, inout, [2.0, 3.0], #=n_outs=# Val(3), Val(1)) + jac = Enzyme.jacobian(Reverse, inout, [2.0, 3.0], n_outs=Val((3,)), chunk=Val(1))[1] @test size(jac) == (3, 2) @test jac ≈ [ 0.0 1.0; 4.0 0.0; 12.0 0.0] - jac = Enzyme.jacobian(Forward, inout, [2.0, 3.0], Val(1)) + jac = Enzyme.jacobian(Forward, inout, [2.0, 3.0], chunk=Val(1))[1] @test size(jac) == (3, 2) @test jac ≈ [ 0.0 1.0; 4.0 0.0; 12.0 0.0] - @test jac == Enzyme.jacobian(Forward, inout, [2.0, 3.0]) + @test jac == Enzyme.jacobian(Forward, inout, [2.0, 3.0])[1] - jac = Enzyme.jacobian(Reverse, inout, [2.0, 3.0], #=n_outs=# Val(3), Val(2)) + jac = Enzyme.jacobian(Reverse, inout, [2.0, 3.0], n_outs=Val((3,)), chunk=Val(2))[1] @test size(jac) == (3, 2) @test jac ≈ [ 0.0 1.0; 4.0 0.0; 12.0 0.0] - jac = Enzyme.jacobian(Forward, inout, [2.0, 3.0], Val(2)) + jac = Enzyme.jacobian(Forward, inout, [2.0, 3.0], chunk=Val(2))[1] @test size(jac) == (3, 2) @test jac ≈ [ 0.0 1.0; 4.0 0.0; @@ -3286,13 +3288,13 @@ end utmp .= A*x[2:end] .+ x[1] end - J_r_1(A, x) = Enzyme.jacobian(Reverse, θ -> f_test_1(A, θ), x, Val(5)) - J_r_2(A, x) = Enzyme.jacobian(Reverse, θ -> f_test_2(A, θ), x, Val(5)) - J_r_3(u, A, x) = Enzyme.jacobian(Reverse, θ -> f_test_3!(u, A, θ), x, Val(5)) + J_r_1(A, x) = Enzyme.jacobian(Reverse, θ -> f_test_1(A, θ), x, n_outs=Val((5,)))[1] + J_r_2(A, x) = Enzyme.jacobian(Reverse, θ -> f_test_2(A, θ), x, n_outs=Val((5,)))[1] + J_r_3(u, A, x) = Enzyme.jacobian(Reverse, θ -> f_test_3!(u, A, θ), x, n_outs=Val((5,)))[1] - J_f_1(A, x) = Enzyme.jacobian(Forward, Const(θ -> f_test_1(A, θ)), x) - J_f_2(A, x) = Enzyme.jacobian(Forward, Const(θ -> f_test_2(A, θ)), x) - J_f_3(u, A, x) = Enzyme.jacobian(Forward, Const(θ -> f_test_3!(u, A, θ)), x) + J_f_1(A, x) = Enzyme.jacobian(Forward, Const(θ -> f_test_1(A, θ)), x)[1] + J_f_2(A, x) = Enzyme.jacobian(Forward, Const(θ -> f_test_2(A, θ)), x)[1] + J_f_3(u, A, x) = Enzyme.jacobian(Forward, Const(θ -> f_test_3!(u, A, θ)), x)[1] x = ones(6) A = Matrix{Float64}(LinearAlgebra.I, 5, 5) @@ -3351,7 +3353,7 @@ end dry = zeros(2) function foo(y, dy, x, dx) - autodiff_deferred(Reverse, speelpenning, Const, Duplicated(y, dy), Duplicated(x, dx)) + autodiff(Reverse, speelpenning, Const, Duplicated(y, dy), Duplicated(x, dx)) return nothing end @@ -3776,8 +3778,8 @@ end @testset "Constant Complex return" begin vec = [0.5] - @test Enzyme.gradient(Enzyme.Reverse, fexpandempty, vec)[1] ≈ 1.0 - @test Enzyme.gradient(Enzyme.Forward, fexpandempty, vec)[1] ≈ 1.0 + @test Enzyme.gradient(Enzyme.Reverse, fexpandempty, vec)[1] ≈ [1.0] + @test Enzyme.gradient(Enzyme.Forward, fexpandempty, vec)[1] ≈ [1.0] end const CUmemoryPool2 = Ptr{Float64} @@ -3924,10 +3926,10 @@ const objective3 = params -> mixture_loglikelihood3(params, data) -13.935687326484112, -38.00044665702692, 12.87712891527131] - @test expected ≈ Enzyme.gradient(Reverse, objective1, params0) + @test expected ≈ Enzyme.gradient(Reverse, objective1, params0)[1] # objective2 fails from runtime activity requirements - # @test expected ≈ Enzyme.gradient(Reverse, objective2, params0) - @test expected ≈ Enzyme.gradient(Reverse, objective3, params0) + # @test expected ≈ Enzyme.gradient(Reverse, objective2, params0)[1] + @test expected ≈ Enzyme.gradient(Reverse, objective3, params0)[1] end struct HarmonicAngle