Merge pull request #59 from JuliaDiff/ox/comp

Introduce Composite for structured deriviatives

Author: oxinabox

src/ChainRulesCore.jl

@@ -1,15 +1,25 @@
 module ChainRulesCore
-using Base.Broadcast: materialize, materialize!, broadcasted, Broadcasted, broadcastable
+using Base.Broadcast: broadcasted, Broadcasted, broadcastable, materialize, materialize!
 
 export frule, rrule
-export wirtinger_conjugate, wirtinger_primal, refine_differential
+export refine_differential, wirtinger_conjugate, wirtinger_primal
 export @scalar_rule, @thunk
 export extern, store!, unthunk
-export Wirtinger, Zero, One, DoesNotExist, Thunk, InplaceableThunk
+export Composite, DoesNotExist, InplaceableThunk, One, Thunk, Wirtinger, Zero
 export NO_FIELDS
 
-include("differentials.jl")
+include("compat.jl")
+
+include("differentials/abstract_differential.jl")
+include("differentials/wirtinger.jl")
+include("differentials/zero.jl")
+include("differentials/does_not_exist.jl")
+include("differentials/one.jl")
+include("differentials/thunks.jl")
+include("differentials/composite.jl")
+
 include("differential_arithmetic.jl")
+
 include("operations.jl")
 include("rules.jl")
 include("rule_definition_tools.jl")

src/compat.jl

@@ -0,0 +1,3 @@
+if VERSION < v"1.2"
+    Base.getproperty(x::Tuple, f::Int) = getfield(x, f)
+end

src/differential_arithmetic.jl

@@ -7,7 +7,8 @@ subtypes, as we know the full set that might be encountered.
 Thus we can avoid any ambiguities.
 
 Notice:
-    The precidence goes: (:Wirtinger, :Zero, :DoesNotExist, :One, :AbstractThunk, :Any)
+    The precedence goes:
+    `Wirtinger, Zero, DoesNotExist, One, AbstractThunk, Composite, Any`
     Thus each of the @eval loops creating definitions of + and *
     defines the combination this type with all types of  lower precidence.
     This means each eval loops is 1 item smaller than the previous.
@@ -87,3 +88,25 @@ for T in (:Any,)
     @eval Base.:*(a::AbstractThunk, b::$T) = unthunk(a) * b
     @eval Base.:*(a::$T, b::AbstractThunk) = a * unthunk(b)
 end
+
+################## Composite ##############################################################
+
+# We intentionally do not define, `Base.*(::Composite, ::Composite)` as that is not meaningful
+# In general one doesn't have to represent multiplications of 2 differentials
+# Only of a differential and a scaling factor (generally `Real`)
+Base.:*(s::Any, comp::Composite) = map(x->s*x, comp)
+Base.:*(comp::Composite, s::Any) = map(x->x*s, comp)
+
+
+function Base.:+(a::Composite{P}, b::Composite{P}) where P
+    data = elementwise_add(backing(a), backing(b))
+    return Composite{P, typeof(data)}(data)
+end
+function Base.:+(a::P, d::Composite{P}) where P
+    try
+        return construct(P, elementwise_add(backing(a), backing(d)))
+    catch err
+        throw(PrimalAdditionFailedException(a, d, err))
+    end
+end
+Base.:+(a::Composite{P}, b::P) where P = b + a

src/differentials.jl

@@ -149,6 +149,7 @@ Base.iterate(::One, ::Any) = nothing
 #####
 ##### `AbstractThunk
 #####
+
 abstract type AbstractThunk <: AbstractDifferential end
 
 Base.Broadcast.broadcastable(x::AbstractThunk) = broadcastable(extern(x))
@@ -237,8 +238,17 @@ macro thunk(body)
     return :(Thunk($(esc(func))))
 end
 
+"""
+    unthunk(x)
+
+`unthunk` removes 1 layer of thunking from an `AbstractThunk`,
+and on all other types is the `identity` function.
+"""
+unthunk(x) = x
+unthunk(x::Thunk) = x()
+
 # have to define this here after `@thunk` and `Thunk` is defined
-Base.conj(x::AbstractThunk) = @thunk(conj(extern(x)))
+Base.conj(x::AbstractThunk) = @thunk(conj(unthunk(x)))
 
 (x::Thunk)() = x.f()
 @inline unthunk(x::Thunk) = x()
@@ -284,6 +294,73 @@ function itself, when that function is not a closure.
 """
 const NO_FIELDS = DoesNotExist()
 
+
+"""
+    Composite{P, T} <: AbstractDifferential
+
+This type represents the differential for a `struct`/`NamedTuple`, or `Tuple`.
+`P` is the the corresponding primal type that this is a differential for.
+
+`Composite{P}` should have fields (technically properties), that match to a subset of the
+fields of the primal type; and each should be a differential type matching to the primal
+type of that field.
+Fields of the P that are not present in the Composite are treated as `Zero`.
+
+`T` is an implementation detail representing the backing data structure.
+For Tuple it will be a Tuple, and for everything else it will be a `NamedTuple`.
+It should not be passed in by user.
+"""
+struct Composite{P, T} <: AbstractDifferential
+    # Note: If T is a Tuple, then P is also a Tuple
+    # (but potentially a different one, as it doesn't contain differentials)
+    backing::T
+end
+
+function Composite{P}(; kwargs...) where P
+    backing = (; kwargs...)  # construct as NamedTuple
+    return Composite{P, typeof(backing)}(backing)
+end
+
+function Composite{P}(args...) where P
+    return Composite{P, typeof(args)}(args)
+end
+
+function Base.show(io::IO, comp::Composite{P}) where P
+    print(io, "Composite{")
+    show(io, P)
+    print(io, "}")
+    # allow Tuple or NamedTuple `show` to do the rendering of brackets etc
+    show(io, backing(comp))
+end
+
+Base.convert(::Type{<:NamedTuple}, comp::Composite{<:Any, <:NamedTuple}) = backing(comp)
+Base.convert(::Type{<:Tuple}, comp::Composite{<:Any, <:Tuple}) = backing(comp)
+
+Base.getindex(comp::Composite, idx) = getindex(backing(comp), idx)
+Base.getproperty(comp::Composite, idx::Int) = getproperty(backing(comp), idx)  # for Tuple
+Base.getproperty(comp::Composite, idx::Symbol) = getproperty(backing(comp), idx)
+Base.propertynames(comp::Composite) = propertynames(backing(comp))
+
+Base.iterate(comp::Composite, args...) = iterate(backing(comp), args...)
+Base.length(comp::Composite) = length(backing(comp))
+Base.eltype(::Type{<:Composite{<:Any, T}}) where T = eltype(T)
+
+function Base.map(f, comp::Composite{P, <:Tuple}) where P
+    vals::Tuple = map(f, backing(comp))
+    return Composite{P, typeof(vals)}(vals)
+end
+function Base.map(f, comp::Composite{P, <:NamedTuple{L}}) where{P, L}
+    vals = map(f, Tuple(backing(comp)))
+    named_vals = NamedTuple{L, typeof(vals)}(vals)
+    return Composite{P, typeof(named_vals)}(named_vals)
+end
+
+Base.conj(comp::Composite) = map(conj, comp)
+
+extern(comp::Composite) = backing(map(extern, comp))  # gives a NamedTuple or Tuple
+
+#==============================================================================#
+
 """
     refine_differential(𝒟::Type, der)
 

src/differentials/abstract_differential.jl

@@ -0,0 +1,52 @@
+#####
+##### `AbstractDifferential`
+#####
+
+"""
+The subtypes of `AbstractDifferential` define a custom \"algebra\" for chain
+rule evaluation that attempts to factor various features like complex derivative
+support, broadcast fusion, zero-elision, etc. into nicely separated parts.
+
+All subtypes of `AbstractDifferential` implement the following operations:
+
+`+(a, b)`: linearly combine differential `a` and differential `b`
+
+`*(a, b)`: multiply the differential `a` by the differential `b`
+
+`Base.conj(x)`: complex conjugate of the differential `x`
+
+`extern(x)`: convert `x` into an appropriate non-`AbstractDifferential` type for
+use outside of `ChainContext`.
+
+Valid arguments to these operations are `T` where `T<:AbstractDifferential`, or
+where `T` has proper `+` and `*` implementations.
+
+Additionally, all subtypes of `AbstractDifferential` support `Base.iterate` and
+`Base.Broadcast.broadcastable(x)`.
+"""
+abstract type AbstractDifferential end
+
+Base.:+(x::AbstractDifferential) = x
+
+"""
+    extern(x)
+
+Return `x` converted to an appropriate non-`AbstractDifferential` type, for use
+with external packages that might not handle `AbstractDifferential` types.
+
+Note that this function may return an alias (not necessarily a copy) to data
+wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
+"""
+@inline extern(x) = x
+
+@inline Base.conj(x::AbstractDifferential) = x
+
+"""
+    refine_differential(𝒟::Type, der)
+
+Converts, if required, a differential object `der`
+(e.g. a `Number`, `AbstractDifferential`, `Matrix`, etc.),
+to another  differential that is more suited for the domain given by the type 𝒟.
+Often this will behave as the identity function on `der`.
+"""
+refine_differential(::Any, der) = der  # most of the time leave it alone.