Add more modes

Author: tansongchen

src/derivative.jl

@@ -1,5 +1,5 @@
 
-export derivative
+export derivative, derivative!
 
 """
     derivative(f, x, order::Int64)
@@ -24,44 +24,62 @@ Batch mode derivative / directional derivative calculations, where each column o
 """
 function derivative end
 
+"""
+    derivative!(result, f, x, l, ::Val{N})
+    derivative!(result, f!, y, x, l, ::Val{N})
+
+In-place derivative calculation APIs. `result` is expected to be pre-allocated and have the same shape as `y`.
+"""
+function derivative! end
+
 # Convenience wrappers for converting orders to value types
 # and forward work to core APIs
 
-@inline function derivative(f, x, order::Int64)
-    derivative(f, x, Val{order + 1}())
-end
-
-@inline function derivative(f, x, l, order::Int64)
-    derivative(f, x, l, Val{order + 1}())
-end
+@inline derivative(f, x, order::Int64) = derivative(f, x, one(eltype(x)), order)
+@inline derivative(f, x, l, order::Int64) = derivative(f, x, l, Val{order + 1}())
 
 # Core APIs
 
 # Added to help Zygote infer types
-make_taylor(t0::T, t1::S, ::Val{N}) where {T, S, N} = TaylorScalar{T, N}(t0, convert(T, t1))
+@inline function make_taylor(x::T, l::S, ::Val{N}) where {T <: TN, S <: TN, N}
+    TaylorScalar{T, N}(x, convert(T, l))
+end
 
-@inline function derivative(f, x::T, ::Val{N}) where {T <: TN, N}
-    t = TaylorScalar{T, N}(x, one(x))
-    return extract_derivative(f(t), N)
+@inline function make_taylor(x::AbstractArray{T}, l, vN::Val{N}) where {T <: TN, N}
+    broadcast(make_taylor, x, l, vN)
 end
 
-@inline function derivative(f, x::AbstractVector{T}, l::AbstractVector{S},
-        vN::Val{N}) where {T <: TN, S <: TN, N}
-    t = map((t0, t1) -> make_taylor(t0, t1, vN), x, l)
-    # equivalent to map(TaylorScalar{T, N}, x, l)
+# Out-of-place function, out-of-place derivative
+@inline function derivative(f, x, l, vN::Val{N}) where {N}
+    t = make_taylor(x, l, vN)
     return extract_derivative(f(t), N)
 end
 
-# shorthand notations for matrices
+# Below three advanced APIs do not have convenience wrappers
+
+# In-place function, out-of-place derivative
+@inline function derivative(f!, y::AbstractArray{T}, x, l, vN::Val{N}) where {T, N}
+    s = similar(y, TaylorScalar{T, N})
+    t = make_taylor(x, l, vN)
+    f!(s, t)
+    map!(primal, y, s)
+    return extract_derivative(s, N)
+end
 
-@inline function derivative(f, x::AbstractMatrix{T}, vN::Val{N}) where {T <: TN, N}
-    size(x)[1] != 1 && @warn "x is not a row vector."
-    t = make_taylor.(x, one(T), vN)
-    return extract_derivative.(f(t), N)
+# Out-of-place function, in-place derivative
+@inline function derivative!(result, f, x, l, vN::Val{N}) where {N}
+    t = make_taylor(x, l, vN)
+    s = f(t)
+    extract_derivative!(result, s, N)
+    return result
 end
 
-@inline function derivative(f, x::AbstractMatrix{T}, l::AbstractVector{S},
-        vN::Val{N}) where {T <: TN, S <: TN, N}
-    t = make_taylor.(x, l, vN)
-    return extract_derivative.(f(t), N)
+# In-place function, in-place derivative
+@inline function derivative!(result, f!, y::AbstractArray{T}, x, l, vN::Val{N}) where {T, N}
+    s = similar(y, TaylorScalar{T, N})
+    t = make_taylor(x, l, vN)
+    f!(s, t)
+    map!(primal, y, s)
+    extract_derivative!(result, s, N)
+    return result
 end

src/scalar.jl

@@ -62,6 +62,10 @@ end
     map(t -> extract_derivative(t, i), v)
 end
 @inline extract_derivative(r, i::Integer) = false
+@inline function extract_derivative!(result::AbstractArray, v::AbstractArray{T},
+        i::Integer) where {T <: TaylorScalar}
+    map!(t -> extract_derivative(t, i), result, v)
+end
 @inline primal(t::TaylorScalar) = extract_derivative(t, 1)
 
 @inline zero(::Type{TaylorScalar{T, N}}) where {T, N} = TaylorScalar{T, N}(zero(T))

test/derivative.jl

@@ -1,16 +1,39 @@
 
-@testset "Derivative" begin
+@testset "O-function, O-derivative" begin
     g(x) = x^3
     @test derivative(g, 1.0, 1) ≈ 3
 
     h(x) = x .^ 3
     @test derivative(h, [2.0 3.0], 1) ≈ [12.0 27.0]
+
+    g1(x) = x[1] * x[1] + x[2] * x[2]
+    @test derivative(g1, [1.0, 2.0], [1.0, 0.0], 1) ≈ 2.0
+
+    h1(x) = sum(x, dims = 1)
+    @test derivative(h1, [1.0 2.0; 2.0 3.0], [1.0, 1.0], 1) ≈ [2.0 2.0]
 end
 
-@testset "Directional derivative" begin
-    g(x) = x[1] * x[1] + x[2] * x[2]
-    @test derivative(g, [1.0, 2.0], [1.0, 0.0], 1) ≈ 2.0
+@testset "I-function, O-derivative" begin
+    g!(y, x) = begin
+        y[1] = x * x
+        y[2] = x + 1
+    end
+    x = 2.0
+    y = [0.0, 0.0]
+    @test derivative(g!, y, x, 1.0, Val{2}()) ≈ [4.0, 1.0]
+end
+
+@testset "O-function, I-derivative" begin
+    g(x) = x .^ 2
+    @test derivative!(zeros(2), g, [1.0, 2.0], [1.0, 0.0], Val{2}()) ≈ [2.0, 0.0]
+end
 
-    h(x) = sum(x, dims = 1)
-    @test derivative(h, [1.0 2.0; 2.0 3.0], [1.0, 1.0], 1) ≈ [2.0 2.0]
+@testset "I-function, I-derivative" begin
+    g!(y, x) = begin
+        y[1] = x[1] * x[1]
+        y[2] = x[2] * x[2]
+    end
+    x = [2.0, 3.0]
+    y = [0.0, 0.0]
+    @test derivative!(y, g!, zeros(2), x, [1.0, 0.0], Val{2}()) ≈ [4.0, 0.0]
 end

test/runtests.jl

@@ -3,5 +3,5 @@ using Test
 
 include("primitive.jl")
 include("derivative.jl")
-include("zygote.jl")
+# include("zygote.jl")
 # include("lux.jl")