Reorganize all the files

Author: oxinabox

src/ChainRulesCore.jl

@@ -8,8 +8,14 @@ export extern, store!, unthunk
 export Composite, DoesNotExist, InplaceableThunk, One, Thunk, Wirtinger, Zero
 export NO_FIELDS
 
-include("differentials.jl")
-include("composite_core.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")

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.

src/composite_core.jl renamed to src/differentials/composite.jl

@@ -1,26 +1,68 @@
-struct PrimalAdditionFailedException{P} <: Exception
-    primal::P
-    differential::Composite{P}
-    original::Exception
+"""
+    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 Base.showerror(io::IO, err::PrimalAdditionFailedException{P}) where {P}
-    println(io, "Could not construct $P after addition.")
-    println(io, "This probably means no default constructor is defined.")
-    println(io, "Either define a default constructor")
-    printstyled(io, "$P(", join(propertynames(err.differential), ", "), ")", color=:blue)
-    println(io, "\nor overload")
-    printstyled(io,
-        "ChainRulesCore.construct(::Type{$P}, ::$(typeof(err.differential)))";
-        color=:blue
-    )
-    println(io, "\nor overload")
-    printstyled(io, "Base.:+(::$P, ::$(typeof(err.differential)))"; color=:blue)
-    println(io, "\nOriginal Exception:")
-    printstyled(io, err.original; color=:yellow)
-    println(io)
+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
+
+
 """
     backing(x)
 
@@ -131,3 +173,36 @@ function elementwise_add(a::NamedTuple{an}, b::NamedTuple{bn}) where {an, bn}
         return NamedTuple{names,types}(vals)
     end
 end
+
+
+struct PrimalAdditionFailedException{P} <: Exception
+    primal::P
+    differential::Composite{P}
+    original::Exception
+end
+
+function Base.showerror(io::IO, err::PrimalAdditionFailedException{P}) where {P}
+    println(io, "Could not construct $P after addition.")
+    println(io, "This probably means no default constructor is defined.")
+    println(io, "Either define a default constructor")
+    printstyled(io, "$P(", join(propertynames(err.differential), ", "), ")", color=:blue)
+    println(io, "\nor overload")
+    printstyled(io,
+        "ChainRulesCore.construct(::Type{$P}, ::$(typeof(err.differential)))";
+        color=:blue
+    )
+    println(io, "\nor overload")
+    printstyled(io, "Base.:+(::$P, ::$(typeof(err.differential)))"; color=:blue)
+    println(io, "\nOriginal Exception:")
+    printstyled(io, err.original; color=:yellow)
+    println(io)
+end
+
+"""
+    NO_FIELDS
+
+Constant for the reverse-mode derivative with respect to a structure that has no fields.
+The most notable use for this is for the reverse-mode derivative with respect to the
+function itself, when that function is not a closure.
+"""
+const NO_FIELDS = DoesNotExist()

src/differentials/does_not_exist.jl

@@ -0,0 +1,18 @@
+"""
+    DoesNotExist()
+
+This differential indicates that the derivative Does Not Exist (D.N.E).
+This is not the cast that it is not implemented, but rather that it mathematically
+is not defined.
+"""
+struct DoesNotExist <: AbstractDifferential end
+
+function extern(x::DoesNotExist)
+    throw(ArgumentError("Derivative does not exit. Cannot be converted to an external type."))
+end
+
+Base.Broadcast.broadcastable(::DoesNotExist) = Ref(DoesNotExist())
+
+Base.iterate(x::DoesNotExist) = (x, nothing)
+Base.iterate(::DoesNotExist, ::Any) = nothing
+

src/differentials/one.jl

@@ -0,0 +1,14 @@
+"""
+     One()
+The Differential which is the multiplicative identity.
+Basically, this represents `1`.
+"""
+struct One <: AbstractDifferential end
+
+extern(x::One) = true  # true is a strong 1.
+
+Base.Broadcast.broadcastable(::One) = Ref(One())
+
+Base.iterate(x::One) = (x, nothing)
+Base.iterate(::One, ::Any) = nothing
+

src/differentials/thunks.jl

@@ -0,0 +1,126 @@
+
+abstract type AbstractThunk <: AbstractDifferential end
+
+Base.Broadcast.broadcastable(x::AbstractThunk) = broadcastable(extern(x))
+
+@inline function Base.iterate(x::AbstractThunk)
+    externed = extern(x)
+    element, state = iterate(externed)
+    return element, (externed, state)
+end
+
+@inline function Base.iterate(::AbstractThunk, (externed, state))
+    element, new_state = iterate(externed, state)
+    return element, (externed, new_state)
+end
+
+#####
+##### `Thunk`
+#####
+
+"""
+    Thunk(()->v)
+A thunk is a deferred computation.
+It wraps a zero argument closure that when invoked returns a differential.
+`@thunk(v)` is a macro that expands into `Thunk(()->v)`.
+
+Calling a thunk, calls the wrapped closure.
+`extern`ing thunks applies recursively, it also externs the differial that the closure returns.
+If you do not want that, then simply call the thunk
+
+```
+julia> t = @thunk(@thunk(3))
+Thunk(var"##7#9"())
+
+julia> extern(t)
+3
+
+julia> t()
+Thunk(var"##8#10"())
+
+julia> t()()
+3
+```
+
+### When to `@thunk`?
+When writing `rrule`s (and to a lesser exent `frule`s), it is important to `@thunk`
+appropriately.
+Propagation rule's that return multiple derivatives are not able to do all the computing themselves.
+ By `@thunk`ing the work required for each, they then compute only what is needed.
+
+#### So why not thunk everything?
+`@thunk` creates a closure over the expression, which (effectively) creates a `struct`
+with a field for each variable used in the expression, and call overloaded.
+
+Do not use `@thunk` if this would be equal or more work than actually evaluating the expression itself. Examples being:
+- The expression wrapping something in a `struct`, such as `Adjoint(x)` or `Diagonal(x)`
+- The expression being a constant
+- The expression being itself a `thunk`
+- The expression being from another `rrule` or `frule` (it would be `@thunk`ed if required by the defining rule already)
+"""
+struct Thunk{F} <: AbstractThunk
+    f::F
+end
+
+
+"""
+    @thunk expr
+
+Define a [`Thunk`](@ref) wrapping the `expr`, to lazily defer its evaluation.
+"""
+macro thunk(body)
+    # Basically `:(Thunk(() -> $(esc(body))))` but use the location where it is defined.
+    # so we get useful stack traces if it errors.
+    func = Expr(:->, Expr(:tuple), Expr(:block, __source__, body))
+    return :(Thunk($(esc(func))))
+end
+
+"""
+    unthunk(x)
+
+On `AbstractThunk`s this removes 1 layer of thunking.
+On any other type, it is the identity operation.
+
+In contrast to `extern` this is nonrecursive.
+"""
+@inline unthunk(x) = x
+
+@inline extern(x::AbstractThunk) = extern(unthunk(x))
+
+# have to define this here after `@thunk` and `Thunk` is defined
+Base.conj(x::AbstractThunk) = @thunk(conj(unthunk(x)))
+
+
+(x::Thunk)() = x.f()
+@inline unthunk(x::Thunk) = x()
+
+Base.show(io::IO, x::Thunk) = println(io, "Thunk($(repr(x.f)))")
+
+"""
+    InplaceableThunk(val::Thunk, add!::Function)
+
+A wrapper for a `Thunk`, that allows it to define an inplace `add!` function,
+which is used internally in `accumulate!(Δ, ::InplaceableThunk)`.
+
+`add!` should be defined such that: `ithunk.add!(Δ) = Δ .+= ithunk.val`
+but it should do this more efficently than simply doing this directly.
+(Otherwise one can just use a normal `Thunk`).
+
+Most operations on an `InplaceableThunk` treat it just like a normal `Thunk`;
+and destroy its inplacability.
+"""
+struct InplaceableThunk{T<:Thunk, F} <: AbstractThunk
+    val::T
+    add!::F
+end
+
+unthunk(x::InplaceableThunk) = unthunk(x.val)
+(x::InplaceableThunk)() = unthunk(x)
+
+function Base.show(io::IO, x::InplaceableThunk)
+    println(io, "InplaceableThunk($(repr(x.val)), $(repr(x.add!)))")
+end
+
+# The real reason we have this:
+accumulate!(Δ, ∂::InplaceableThunk) = ∂.add!(Δ)
+store!(Δ, ∂::InplaceableThunk) = ∂.add!((Δ.*=false))  # zero it, then add to it.