See also: mim::plug::autodiff

This plugin provides reverse-mode automatic differentiation: autodiff.ad f yields the augmented function f' = λ args. (f args, f*) that returns the result together with its pullback. The actual differentiation is performed by the mim::plug::autodiff::Eval phase.

Dependencies

plugin mem;
import compile;
plugin core; // for derivatives

Types

Tangent

The tangent type of a type T (currently the identity).

axm Tangent: * → *, normalize_Tangent;

Operations

AD

Computes the augmented type of a function type: (T → U) => (T → U × (U → T)).

axm AD: * → *, normalize_AD;

ad

Yields the augmented term of a closed term (function, operator, higher-order argument, registered axiom, etc.). The augmented term f' returns the result together with the pullback: autodiff f = f' = λ args. (f args, f*)

axm ad: {T: *} → T → AD T, normalize_ad;

zero

Represents universal zero such that (zero T) +_T t = t.

axm zero: [T: *] → T, normalize_zero;

add

A universal addition that consumes zeros and defaults to normal addition for scalar types. It lifts addition over structured types and special-cases types that do not allow addition.

axm add: {T: *} → [T, T] → T, normalize_add;

sum

Performs autodiff.add over a list of terms.

axm sum: [n: Nat, T: *] → «n; T» → T, normalize_sum;

Phases

zero_repl

Replaces autodiff.zero T with the zero element of T, where one exists.

axm zero_repl: compile.Phase;

eval

The heart of AD: replaces an autodiff.ad call with the differentiated function (mim::plug::autodiff::phase::Eval).

axm eval: compile.Phase;

Registered translations

In this section, we define translations for axioms of other plugins. This would best be done using a register mechanism in a third plugin or at least in a separate file.

The general concept is that a call to an axiom is replaced with a call to the augmented axiom. The augmented axiom needs a wrapper for meta arguments (a higher-order function). Appropriate cps2ds wrappers are introduced to handle that the augmented axioms are in CPS whereas the original axioms are in direct style. Example:

mul' => args → result*pullback
call: r    = mul  (m,w) (a,b)
res : r,r* = mul' (m,w) (a,b)

The types (with Int for (Int w)) are:

mul : [m:Nat,w:Nat] → [a:Int,b:Int] → Int
r   : Int
r*  : cn[Int,cn[Int,Int]]

The pullback has to be in cps for compliance.

mul* := λ s. (s*b,s*a)
mul'_cps : [m:Nat,w:Nat] → cn[[Int,Int],cn[Int,   cn[Int,cn[Int,Int]]]]
r,r* = (cps2ds (mul'_cps (m,w))) (a,b)

The pullback is the derivative with respect to the input (weighted with the output tangent). For arithmetic operations, s is simply multiplied onto each input tangent: ∂_i f(x1,...,xn) * s It follows that the applied partial pullback needs to be: sum x_i*(∂_i f(x1,...,xn) * s) = sum x_i*(•) with • as the formula from above. This is a direct result of the chain-rule composition with the partial pullback of a tuple. The tuple pullback transports the partial pullbacks of the operands and handles the sums. By its nature the pullback of a tuple needs to be a sum.

diff

The derivative of the axiom plugin.tag.sub is the annex autodiff.diff.plugin_tag_sub, i.e. the axiom's name with all .s replaced by _s; this is how mim::plug::autodiff::phase::Eval looks a derivative up. Being annexes, they are roots of the World and hence survive until mim::plug::autodiff::phase::Eval has consumed them. But - unlike externals - they are invisible to the backends and a compile.unload "autodiff" nukes them along with all the other annexes of this plugin.

mod diff {

autodiff.diff.core_icmp_xYgLE

The derivative of core.icmp.xYgLE (eq).

The comparison pullback exists formally but is not used.

anx fun core_icmp_xYgLE {s: Nat} (ab: «2; Idx s»)@tt: [Bool, Cn [Bool, Cn «2; Idx s»]] =

return (core.icmp.sle ab, fn (_: Bool)@tt: «2; Idx s» = return ‹2; 0:(Idx s)›);

autodiff.diff.core_wrap_add

The derivative of core.wrap.add:

s ↦ (s, s)

anx fun core_wrap_add (m: Nat) {s: Nat} (ab: «2; Idx s»)@tt: [Idx s, Cn [Idx s, Cn «2; Idx s»]] =

return (core.wrap.add m ab, fn (i: Idx s)@tt: «2; Idx s» = return ‹2; i›);

autodiff.diff.core_wrap_mul

The derivative of core.wrap.mul:

s ↦ (s*b, s*a)

    anx fun core_wrap_mul (m: Nat) {s: Nat} (a b: Idx s) as ab@tt: [Idx s, Cn [Idx s, Cn «2; Idx s»]] =
        return (core.wrap.mul m ab, fn (i: Idx s)@tt: «2; Idx s» = return (core.wrap.mul m (i, b), core.wrap.mul m (i, a)));
}