=WIP Derivative wrt function

=Make frule wrt self and rrule wrt self different [WIP

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.2.1-DEV"
+version = "v0.3.0"
 
 [compat]
 julia = "^1.0"

src/ChainRulesCore.jl

@@ -4,6 +4,7 @@ using Base.Broadcast: materialize, materialize!, broadcasted, Broadcasted, broad
 export AbstractRule, Rule, frule, rrule
 export @scalar_rule, @thunk
 export extern, cast, store!, Wirtinger, Zero, One, Casted, DNE, Thunk, DNERule
+export NO_FIELDS_RULE, ZERO_RULE
 
 include("differentials.jl")
 include("differential_arithmetic.jl")

src/rule_definition_tools.jl

@@ -14,15 +14,17 @@ methods for `frule` and `rrule`:
     function ChainRulesCore.frule(::typeof(f), x₁::Number, x₂::Number, ...)
         Ω = f(x₁, x₂, ...)
         \$(statement₁, statement₂, ...)
-        return Ω, (Rule((Δx₁, Δx₂, ...) -> ∂f₁_∂x₁ * Δx₁ + ∂f₁_∂x₂ * Δx₂ + ...),
+        return Ω, (ZERO_RULE,
+                   Rule((Δx₁, Δx₂, ...) -> ∂f₁_∂x₁ * Δx₁ + ∂f₁_∂x₂ * Δx₂ + ...),
                    Rule((Δx₁, Δx₂, ...) -> ∂f₂_∂x₁ * Δx₁ + ∂f₂_∂x₂ * Δx₂ + ...),
                    ...)
     end
 
     function ChainRulesCore.rrule(::typeof(f), x₁::Number, x₂::Number, ...)
         Ω = f(x₁, x₂, ...)
         \$(statement₁, statement₂, ...)
-        return Ω, (Rule((ΔΩ₁, ΔΩ₂, ...) -> ∂f₁_∂x₁ * ΔΩ₁ + ∂f₂_∂x₁ * ΔΩ₂ + ...),
+        return Ω, (NO_FIELDS_RULE,
+                   Rule((ΔΩ₁, ΔΩ₂, ...) -> ∂f₁_∂x₁ * ΔΩ₁ + ∂f₂_∂x₁ * ΔΩ₂ + ...),
                    Rule((ΔΩ₁, ΔΩ₂, ...) -> ∂f₁_∂x₂ * ΔΩ₁ + ∂f₂_∂x₂ * ΔΩ₂ + ...),
                    ...)
     end
@@ -34,11 +36,16 @@ Constraints may also be explicitly be provided to override the `Number` constrai
 e.g. `f(x₁::Complex, x₂)`, which will constrain `x₁` to `Complex` and `x₂` to
 `Number`.
 
-Note that the result of `f(x₁, x₂, ...)` is automatically bound to `Ω`. This
+At present this does not support defining rules for closures/functors.
+This the first returned rule, representing the derivative with respect to the
+function itself, is always the `NO_FIELDS_RULE` (reverse-mode),
+or `ZERO_RULE` (forward-mode).
+
+The result of `f(x₁, x₂, ...)` is automatically bound to `Ω`. This
 allows the primal result to be conveniently referenced (as `Ω`) within the
 derivative/setup expressions.
 
-Note that the `@setup` argument can be elided if no setup code is need. In other
+The `@setup` argument can be elided if no setup code is need. In other
 words:
 
     @scalar_rule(f(x₁, x₂, ...),
@@ -92,9 +99,18 @@ macro scalar_rule(call, maybe_setup, partials...)
         forward_rules = Any[rule_from_partials(inputs[1], partial) for partial in partials]
         reverse_rules = Any[rule_from_partials(inputs[1], partials...)]
     end
-    forward_rules = length(forward_rules) == 1 ? forward_rules[1] : Expr(:tuple, forward_rules...)
-    reverse_rules = length(reverse_rules) == 1 ? reverse_rules[1] : Expr(:tuple, reverse_rules...)
+
+    # First pseudo-partial is derivative WRT function itself.  Since this macro does not
+    # support closures, it is just the empty NamedTuple
+    forward_rules = Expr(:tuple, ZERO_RULE, forward_rules...)
+    reverse_rules = Expr(:tuple, NO_FIELDS_RULE, reverse_rules...)
     return quote
+        if fieldcount(typeof($f)) > 0
+            throw(ArgumentError(
+                "@scalar_rule cannot be used on closures/functors (such as $f)"
+            ))
+        end
+
         function ChainRulesCore.frule(::typeof($f), $(inputs...))
             $(esc(:Ω)) = $call
             $(setup_stmts...)

src/rules.jl

@@ -1,3 +1,242 @@
+"""
+Subtypes of `AbstractRule` are types which represent the primitive derivative
+propagation "rules" that can be composed to implement forward- and reverse-mode
+automatic differentiation.
+
+More specifically, a `rule::AbstractRule` is a callable Julia object generally
+obtained via calling [`frule`](@ref) or [`rrule`](@ref). Such rules accept
+differential values as input, evaluate the chain rule using internally stored/
+computed partial derivatives to produce a single differential value, then
+return that calculated differential value.
+
+For example:
+
+```jldoctest
+julia> using ChainRulesCore: frule, rrule, AbstractRule
+
+julia> x, y = rand(2);
+
+julia> h, dh = frule(hypot, x, y);
+
+julia> h == hypot(x, y)
+true
+
+julia> isa(dh, AbstractRule)
+true
+
+julia> Δx, Δy = rand(2);
+
+julia> dh(Δx, Δy) == ((x / h) * Δx + (y / h) * Δy)
+true
+
+julia> h, (dx, dy) = rrule(hypot, x, y);
+
+julia> h == hypot(x, y)
+true
+
+julia> isa(dx, AbstractRule) && isa(dy, AbstractRule)
+true
+
+julia> Δh = rand();
+
+julia> dx(Δh) == (x / h) * Δh
+true
+
+julia> dy(Δh) == (y / h) * Δh
+true
+```
+
+See also: [`frule`](@ref), [`rrule`](@ref), [`Rule`](@ref), [`DNERule`](@ref), [`WirtingerRule`](@ref)
+"""
+abstract type AbstractRule end
+
+# this ensures that consumers don't have to special-case rule destructuring
+Base.iterate(rule::AbstractRule) = (rule, nothing)
+Base.iterate(::AbstractRule, ::Any) = nothing
+
+# This ensures we don't need to check whether the result of `rrule`/`frule` is a tuple
+# in order to get the `i`th rule (assuming it's 1)
+Base.getindex(rule::AbstractRule, i::Integer) = i == 1 ? rule : throw(BoundsError())
+
+"""
+    accumulate(Δ, rule::AbstractRule, args...)
+
+Return `Δ + rule(args...)` evaluated in a manner that supports ChainRulesCore'
+various `AbstractDifferential` types.
+
+This method intended to be customizable for specific rules/input types. For
+example, here is pseudocode to overload `accumulate` w.r.t. a specific forward
+differentiation rule for a given function `f`:
+
+```
+df(x) = # forward differentiation primitive implementation
+
+frule(::typeof(f), x) = (f(x), Rule(df))
+
+accumulate(Δ, rule::Rule{typeof(df)}, x) = # customized `accumulate` implementation
+```
+
+See also: [`accumulate!`](@ref), [`store!`](@ref), [`AbstractRule`](@ref)
+"""
+accumulate(Δ, rule::AbstractRule, args...) = add(Δ, rule(args...))
+
+"""
+    accumulate!(Δ, rule::AbstractRule, args...)
+
+Similar to [`accumulate`](@ref), but compute `Δ + rule(args...)` in-place,
+storing the result in `Δ`.
+
+Note that this function internally calls `Base.Broadcast.materialize!(Δ, ...)`.
+
+See also: [`accumulate`](@ref), [`store!`](@ref), [`AbstractRule`](@ref)
+"""
+function accumulate!(Δ, rule::AbstractRule, args...)
+    return materialize!(Δ, broadcastable(add(cast(Δ), rule(args...))))
+end
+
+accumulate!(Δ::Number, rule::AbstractRule, args...) = accumulate(Δ, rule, args...)
+
+"""
+    store!(Δ, rule::AbstractRule, args...)
+
+Compute `rule(args...)` and store the result in `Δ`, potentially avoiding
+intermediate temporary allocations that might be necessary for alternative
+approaches (e.g. `copyto!(Δ, extern(rule(args...)))`)
+
+Note that this function internally calls `Base.Broadcast.materialize!(Δ, ...)`.
+
+Like [`accumulate`](@ref) and [`accumulate!`](@ref), this function is intended
+to be customizable for specific rules/input types.
+
+See also: [`accumulate`](@ref), [`accumulate!`](@ref), [`AbstractRule`](@ref)
+"""
+store!(Δ, rule::AbstractRule, args...) = materialize!(Δ, broadcastable(rule(args...)))
+
+#####
+##### `Rule`
+#####
+
+Cassette.@context RuleContext
+
+const RULE_CONTEXT = Cassette.disablehooks(RuleContext())
+
+Cassette.overdub(::RuleContext, ::typeof(+), a, b) = add(a, b)
+Cassette.overdub(::RuleContext, ::typeof(*), a, b) = mul(a, b)
+
+Cassette.overdub(::RuleContext, ::typeof(add), a, b) = add(a, b)
+Cassette.overdub(::RuleContext, ::typeof(mul), a, b) = mul(a, b)
+
+"""
+    Rule(propation_function[, updating_function])
+
+Return a `Rule` that wraps the given `propation_function`. It is assumed that
+`propation_function` is a callable object whose arguments are differential
+values, and whose output is a single differential value calculated by applying
+internally stored/computed partial derivatives to the input differential
+values.
+
+If an updating function is provided, it is assumed to have the signature `u(Δ, xs...)`
+and to store the result of the propagation function applied to the arguments `xs` into
+`Δ` in-place, returning `Δ`.
+
+For example:
+
+```
+frule(::typeof(*), x, y) = x * y, Rule((Δx, Δy) -> Δx * y + x * Δy)
+
+rrule(::typeof(*), x, y) = x * y, (Rule(ΔΩ -> ΔΩ * y'), Rule(ΔΩ -> x' * ΔΩ))
+```
+
+See also: [`frule`](@ref), [`rrule`](@ref), [`accumulate`](@ref), [`accumulate!`](@ref), [`store!`](@ref)
+"""
+struct Rule{F,U<:Union{Function,Nothing}} <: AbstractRule
+    f::F
+    u::U
+end
+
+# NOTE: Using `Core.Typeof` instead of `typeof` here so that if we define a rule for some
+# constructor based on a `UnionAll`, we get `Rule{Type{Thing}}` instead of `Rule{UnionAll}`
+Rule(f) = Rule{Core.Typeof(f),Nothing}(f, nothing)
+
+(rule::Rule{F})(args...) where {F} = Cassette.overdub(RULE_CONTEXT, rule.f, args...)
+
+# Specialized accumulation
+# TODO: Does this need to be overdubbed in the rule context?
+accumulate!(Δ, rule::Rule{F,U}, args...) where {F,U<:Function} = rule.u(Δ, args...)
+
+
+"""
+    NO_FIELDS_RULE
+
+Constant for the rule for the derivative with respect to 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.
+The rule returns an empty `NamedTuple` for all inputs.
+"""
+const NO_FIELDS_RULE = Rule((args...)->NamedTuple())
+
+"""
+    ZERO_RULE
+
+This is a rule that returns `Zero()` regardless of input.
+The most notable use for this is for the forward-mode derivative with respect to the
+function itself, when that function is not a closure.
+"""
+const ZERO_RULE = Rule((args...)->Zero())
+
+
+
+#####
+##### `DNERule`
+#####
+
+"""
+    DNERule(args...)
+
+Construct a `DNERule` object, which is an `AbstractRule` that signifies that the
+current function is not differentiable with respect to a particular parameter.
+**DNE** is an abbreviation for Does Not Exist.
+"""
+struct DNERule <: AbstractRule end
+
+DNERule(args...) = DNE()
+
+#####
+##### `WirtingerRule`
+#####
+
+"""
+    WirtingerRule(primal::AbstractRule, conjugate::AbstractRule)
+
+Construct a `WirtingerRule` object, which is an `AbstractRule` that consists of
+an `AbstractRule` for both the primal derivative ``∂/∂x`` and the conjugate
+derivative ``∂/∂x̅``. If the domain `𝒟` of the function might be real, consider
+calling `AbstractRule(𝒟, primal, conjugate)` instead, to make use of a more
+efficient representation wherever possible.
+"""
+struct WirtingerRule{P<:AbstractRule,C<:AbstractRule} <: AbstractRule
+    primal::P
+    conjugate::C
+end
+
+function (rule::WirtingerRule)(args...)
+    return Wirtinger(rule.primal(args...), rule.conjugate(args...))
+end
+
+"""
+    AbstractRule(𝒟::Type, primal::AbstractRule, conjugate::AbstractRule)
+
+Return a `Rule` evaluating to `primal(Δ) + conjugate(Δ)` if `𝒟 <: Real`,
+otherwise return `WirtingerRule(P, C)`.
+"""
+function AbstractRule(𝒟::Type, primal::AbstractRule, conjugate::AbstractRule)
+    if 𝒟 <: Real || eltype(𝒟) <: Real
+        return Rule((args...) -> add(primal(args...), conjugate(args...)))
+    else
+        return WirtingerRule(primal, conjugate)
+    end
+end
+
 #####
 ##### `frule`/`rrule`
 #####