The automatic differentiation Plugin

Table of Contents

• Dependencies
• Types
  • %autodiff.Tangent
• Operations
  • %autodiff.AD
  • %autodiff.ad
  • %autodiff.zero
  • %autodiff.add
  • %autodiff.sum
• Phases
• Registered translations
  • %core.icmp.xYgLE (eq)
  • %core.wrap.add
  • %core.wrap.mul

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

%autodiff.Tangent

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

axm %autodiff.Tangent: * → *, normalize_Tangent;

Operations

%autodiff.AD

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

axm %autodiff.AD: * → *, normalize_AD;

%autodiff.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 %autodiff.ad: {T: *} → T → %autodiff.AD T, normalize_ad;

%autodiff.zero

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

axm %autodiff.zero: [T: *] → T, normalize_zero;

%autodiff.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.

TODO: how do we handle summations that need memory? (grab current memory?)

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

%autodiff.sum

Performs %autodiff.add over a list of terms.

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

Phases

axm %autodiff.zero_repl: %compile.Phase;
axm %autodiff.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.

%core.icmp.xYgLE (eq)

The comparison pullback exists formally but is not used.

fun extern internal_diff_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)›);

%core.wrap.add

s ↦ (s, s)

fun extern internal_diff_core_wrap_add {s: Nat} (m: 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›);

%core.wrap.mul

s ↦ (s*b, s*a)

fun extern internal_diff_core_wrap_mul {s: Nat} (m: 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)));