Merge pull request #13 from YingboMa/s/api

User API & DualArray

Author: YingboMa

Project.toml

@@ -7,6 +7,7 @@ version = "0.1.0"
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 

src/ForwardDiff2.jl

@@ -5,7 +5,7 @@ include("custom_dispatch.jl")
 include("tag.jl")
 include("dualarray.jl")
 include("dual_context.jl")
-include("jacobian.jl")
+include("api.jl")
 
 # Experimental
 #include("aosoa.jl")

src/api.jl

@@ -0,0 +1,53 @@
+using StaticArrays: StaticArray, SMatrix, SVector
+using LinearAlgebra: Diagonal, I
+
+extract_diffresult(xs::AbstractArray{<:Number}) = xs
+# need to optimize
+extract_diffresult(xs) = hcat(xs...)'
+function extract_diffresult(xs::StaticVector{<:StaticArray})
+    tup = reduce((x,y)->tuple(x..., y...), map(x->x.data, xs.data))
+    SMatrix{length(xs), length(xs[1])}(tup)
+end
+extract_diffresult(xs::AbstractMatrix{<:Number}) = xs
+extract_diffresult(xs::AbstractVector{<:Number}) = xs'
+
+allpartials(xs) = map(partials, xs)
+
+function seed(v::SVector{N}) where N
+    SMatrix{N,N,eltype(v)}(I)
+end
+
+function seed(v)
+    Matrix(Diagonal(map(one, v)))
+end
+
+function D(f)
+    # grad
+    function deriv(arg::AbstractArray)
+        # always chunk
+        darr = dualrun(()->DualArray(arg, seed(arg)))
+        res = dualrun(()->f(darr))
+        diffres = extract_diffresult(allpartials(res))
+        return diffres
+    end
+    # scalar
+    function deriv(arg)
+        dualrun() do
+            dualized = map(x->Dual(x, one(x)), arg)
+            res = f(dualized)
+            return map(partials, res)
+        end
+    end
+    return deriv
+end
+
+#=
+# scalar case: f: R -> something
+D(sin)(1.0)
+D(x->[x, x^2])(3)
+
+# gradient case: f: R^n -> R
+D(sum)([1,2,3])
+
+# Jacobian case: f: R^n -> R^m
+=#

src/dual_context.jl

@@ -47,11 +47,6 @@ Wirtinger(primal, conjugate) = Wirtinger.(primal, conjugate)
     Core._apply(alternative, (ctx, g), args...)
 end
 
-# this makes `log` work by making throw_complex_domainerror inferable, but not really sure why
-@inline isinteresting(ctx::TaggedCtx, f::typeof(Core.throw), xs) = true
-# add `DualContext` here to avoid ambiguity
-@noinline alternative(ctx::Union{DualContext,TaggedCtx}, f::typeof(Core.throw), arg) = throw(arg)
-
 # actually interesting:
 
 @inline isinteresting(ctx::TaggedCtx, f, a) = anydual(a)
@@ -163,3 +158,13 @@ for pred in BINARY_PREDICATES
         $pred(vx, vy)
     end
 end
+
+
+##### Inference Hacks
+# this makes `log` work by making throw_complex_domainerror inferable, but not really sure why
+@inline isinteresting(ctx::TaggedCtx, f::typeof(Core.throw), xs) = true
+# add `DualContext` here to avoid ambiguity
+@noinline alternative(ctx::Union{DualContext,TaggedCtx}, f::typeof(Core.throw), arg) = throw(arg)
+
+@inline isinteresting(ctx::TaggedCtx, f::typeof(Base.print_to_string), args...) = true
+@noinline alternative(ctx::Union{DualContext,TaggedCtx}, f::typeof(Base.print_to_string), args...) = f(args...)

src/dualarray.jl

@@ -1,10 +1,23 @@
 using StaticArrays: SVector
+
+partial_type(::Dual{T,V,P}) where {T,V,P} = P
 # TODO: Tagging?
 # TODO: Integrate better with SVector. Maybe even use SIMD.jl?
 
-struct DualArray{T,E,M,D<:AbstractArray,I} <: AbstractArray{E,M}
-    data::D
-    DualArray{T,I}(a::AbstractArray{E,N}) where {T,I,E,N} = new{T,E,N-1,typeof(a),I}(a)
+struct DualArray{T,E,M,V<:AbstractArray,D<:AbstractArray} <: AbstractArray{E,M}
+    data::V
+    partials::D
+    function DualArray{T}(v::AbstractArray{E,N}, p::P) where {T,E,N,P<:AbstractArray}
+        # TODO: non-allocating X?
+        X = typeof(similar(p, Base.tail(ntuple(_->0, Val(ndims(P))))))
+        # we need the eltype of `DualArray` to be `Dual{T,E,X}` as opposed to
+        # some kind of `view`, because we can convert `SubArray` to `Array` but
+        # not vise a versa.
+        #
+        # We need that to differentiate through the following code
+        # `(foo(x::AbstractArray{T})::T) where {T} = x[1]`
+        return new{T,Dual{T,E,X},N,typeof(v),typeof(p)}(v, p)
+    end
 end
 
 ###
@@ -18,91 +31,39 @@ function Base.print_array(io::IO, da::DualArray)
     Base.printstyled(io, "Primals:\n", bold=false, color=2)
     prev_params = io isa IOContext ? io.dict : ()
     ioc = IOContext(io, prev_params..., sz)
-    Base.print_array(ioc, value(da))
+    Base.print_array(ioc, data(da))
     Base.println(io)
-    for i=1:npartials(da)
-        Base.printstyled(io,"Partials($i):\n", bold=false, color=3)
-        Base.print_array(ioc, getindex.(partials.(da), i))
-        i !== npartials(da) && Base.println(io)
-    end
+    Base.printstyled(io,"Partials:\n", bold=false, color=3)
+    Base.print_array(ioc, allpartials(da))
     return nothing
 end
 
-DualArray(a::AbstractArray) = DualArray{Nothing,size(a, ndims(a))-1}(a)
-npartials(d::DualArray{T,E,M,D,I}) where {T,E,M,D,I} = I
-tagtype(::Type{<:DualArray{T}}) where {T} = T
-tagtype(::T) where {T<:DualArray} = tagtype(T)
+DualArray(a::AbstractArray, b::AbstractArray) = DualArray{typeof(dualtag())}(a, b)
+npartials(d::DualArray) = size(d.partials, ndims(d.partials))
 data(d::DualArray) = d.data
+allpartials(d::DualArray) = d.partials
 
 ###
 ### Array interface
 ###
 
-Base.eltype(d::DualArray{T,E}) where {T,E} = Dual{T,E,npartials(d)}
-droplast(d::Tuple) = d |> reverse |> Base.tail |> reverse
-Base.size(d::DualArray) = size(data(d)) |> droplast
-Base.size(d::DualArray, i) = i <= ndims(d) ? size(data(d), i) : 1
+#droplast(d::Tuple) = d |> reverse |> Base.tail |> reverse
+Base.size(d::DualArray) = size(data(d))
 Base.IndexStyle(d::DualArray) = Base.IndexStyle(data(d))
-Base.strides(d::DualArray) = strides(data(d)) |> droplast
-Base.similar(d::DualArray{T}, ::Type{S}, dims::Dims) where {T,S} = DualArray{T,npartials(d)}(similar(data(d), S, (dims..., npartials(d)+1)))
+Base.similar(d::DualArray{T}, ::Type{S}, dims::Dims) where {T, S} = DualArray{T}(similar(data(d)), similar(allpartials(d)))
+Base.eachindex(d::DualArray) = eachindex(data(d))
 
-###
-### Broadcast interface
-###
+using StaticArrays
 
-using Base.Broadcast: Broadcasted, BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle
+Base.@propagate_inbounds _slice(A, i...) = @view A[i..., :]
+Base.@propagate_inbounds _slice(A::StaticArray, i...) = A[i..., :]
 
-struct DualStyle{M,T,I,D} <: AbstractArrayStyle{M}
+Base.@propagate_inbounds function Base.getindex(d::DualArray{T}, i::Int...) where {T}
+    return Dual{T}(data(d)[i...], _slice(allpartials(d), i...))
 end
 
-Base.Broadcast.BroadcastStyle(::Type{<:DualArray{T,E,M,D,I}}) where {T,E,M,D,I} = DualStyle{M,T,I,typeof(Base.Broadcast.BroadcastStyle(D))}()
-function Base.similar(bc::Broadcasted{<:DualStyle{M,T,I,D}}, ::Type{E}) where {M,T,I,D,E}
-    if E <: Dual
-        V = valtype(E)
-        arr = DualArray{T,I}(similar(Array{V}, (axes(bc)..., Base.OneTo(I+1)))) # TODO: work with arbitrary array types. Maybe use `ArrayInterface.jl`?
-    else
-        bc′ = convert(Broadcasted{D}, bc)
-        arr = Base.similar(bc′, E)
-    end
-    return arr
-end
-Base.BroadcastStyle(::DualStyle{M,T,I,D}, ::DualStyle{M,T,I,V}) where {M,T,I,D,V} = DualStyle{M,T,I,typeof(Base.BroadcastStyle(D(), V()))}()
-Base.BroadcastStyle(::DualStyle{M,T,I,D}, B::BroadcastStyle) where {M,T,I,D} = DualStyle{M,T,I,typeof(Base.BroadcastStyle(D(), B))}()
-Base.BroadcastStyle(::DualStyle{M,T,I,D}, B::DefaultArrayStyle) where {M,T,I,D} = DualStyle{M,T,I,typeof(Base.BroadcastStyle(D(), B))}()
-
-function value(d::DualArray)
-    n = ndims(d)
-    dd = data(d)
-    return @view dd[ntuple(_ -> Colon(), Val(n))..., 1]
-end
-
-function partials(d::DualArray)
-    n = ndims(d)
-    dd = data(d)
-    return @view dd[ntuple(_ -> Colon(), Val(n))..., 2:end]
-end
-
-Base.eachindex(d::DualArray) = eachindex(@view data(d)[:, 1])
-
-Base.@propagate_inbounds function Base.getindex(d::DualArray, i::Int...)
-    dd = data(d)
-    # TODO: do something different if dd is not Linear index style
-    ii = LinearIndices(size(d))[i...]
-    val   = dd[ii]
-    slice_len = length(d)
-    parts = ntuple(j->dd[j * slice_len + ii], Val(npartials(d)))
-    return Dual(val, SVector(parts))
-end
-
-Base.@propagate_inbounds function Base.setindex!(d::DualArray, dual, i::Int...)
-    dd = data(d)
-    ii = LinearIndices(size(d))[i...]
-    dd[ii] = value(dual)
-
-    slice_len = length(d)
-    ps = partials(dual)
-    for j = 1:npartials(d)
-        dd[j * slice_len + ii] = ps[j]
-    end
+Base.@propagate_inbounds function Base.setindex!(d::DualArray{T}, dual::Dual{T}, i::Int...) where {T}
+    data(d)[i...] = value(dual)
+    allpartials(d)[i..., :] .= partials(dual)
     return dual
 end

src/jacobian.jl

@@ -0,0 +1,13 @@
+using Test
+using ForwardDiff2: D
+using StaticArrays
+
+@testset begin
+    @test D(sin)(1.0) == cos(1.0)
+    @test D(x->[x, x^2])(3) == [1, 6]
+    @test D(sum)([1,2,3]) == ones(3)'
+    @test D(x->@SVector([x[1]^x[2], x[3]^3, x[3]*x[2]*x[1]]))(@SVector[1,2,3.]) === @SMatrix [2.0 0 0; 0 0 27; 6 3 2]
+    @test D(cumsum)(@SVector([1,2,3])) == @SMatrix [1 0 0; 1 1 0; 1 1 1]
+    @test D(cumsum)([1,2,3]) == [1 0 0; 1 1 0; 1 1 1]
+    @test D(x->@SVector([x[1], x[2]]))(@SVector([1,2,3])) === @SMatrix [1 0 0; 0 1 0]
+end

test/api.jl

@@ -1,5 +1,4 @@
 using SafeTestsets
+@time @safetestset "API Tests" begin include("api.jl") end
 @time @safetestset "Nested Differentiation Tests" begin include("nested.jl") end
 @time @safetestset "Dual Tests" begin include("dualtest.jl") end
-@time @safetestset "DualArray Tests" begin include("dualarray.jl") end
-@time @safetestset "Jacobian Tests" begin include("jacobian.jl") end