Reverse-Mode Automatic Differentiation

Implementation of reverse-mode AD (backpropagation) similar to Haskell's ad package. This enables efficient gradient computation for machine learning optimization.

Features

• Dual numbers with forward value and backward gradient
• Reverse-mode differentiation (backpropagation)
• Support for common math operations
• Efficient gradient computation via computational graph

Usage Example

-- Define a function
def f (x : Float) : Float := x * x + 2.0 * x + 1.0

-- Compute gradient at x = 3.0
let grad := diff f 3.0
-- Result: 8.0 (derivative of x² + 2x + 1 at x=3 is 2x + 2 = 8)

structure Hesper.AD.Reverse.TapeEntry : Type

Tape entry for reverse-mode AD. Stores the operation index, parent indices, and local derivatives.

• idx : Nat

  Index of this operation in the tape

• parents : List Nat

  Parent indices (inputs to this operation)

• localGrads : List Float

  Local derivatives with respect to each parent

def Hesper.AD.Reverse.instInhabitedTapeEntry.default : TapeEntry

Equations

• Hesper.AD.Reverse.instInhabitedTapeEntry.default = { idx := default, parents := default, localGrads := default }

instance Hesper.AD.Reverse.instInhabitedTapeEntry : Inhabited TapeEntry

Equations

• Hesper.AD.Reverse.instInhabitedTapeEntry = { default := Hesper.AD.Reverse.instInhabitedTapeEntry.default }

instance Hesper.AD.Reverse.instReprTapeEntry : Repr TapeEntry

Equations

• Hesper.AD.Reverse.instReprTapeEntry = { reprPrec := Hesper.AD.Reverse.instReprTapeEntry.repr }

def Hesper.AD.Reverse.instReprTapeEntry.repr : TapeEntry → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

structure Hesper.AD.Reverse.Tape : Type

Computational tape for reverse-mode AD

• entries : Array TapeEntry

  Entries in the tape (reverse chronological order)

• nextIdx : Nat

  Current index counter

def Hesper.AD.Reverse.instInhabitedTape.default : Tape

Equations

• Hesper.AD.Reverse.instInhabitedTape.default = { entries := default, nextIdx := default }

instance Hesper.AD.Reverse.instInhabitedTape : Inhabited Tape

Equations

• Hesper.AD.Reverse.instInhabitedTape = { default := Hesper.AD.Reverse.instInhabitedTape.default }

def Hesper.AD.Reverse.Tape.empty : Tape

Create an empty tape

Equations

• Hesper.AD.Reverse.Tape.empty = { entries := #[], nextIdx := 1 }

def Hesper.AD.Reverse.Tape.add (tape : Tape) (parents : List Nat) (localGrads : List Float) : Tape × Nat

Add an entry to the tape

Equations

• tape.add parents localGrads = ({ entries := tape.entries.push { idx := tape.nextIdx, parents := parents, localGrads := localGrads }, nextIdx := tape.nextIdx + 1 }, tape.nextIdx)

structure Hesper.AD.Reverse.Dual : Type

Dual number for reverse-mode AD. Contains the primal value and a reference to the tape.

• primal : Float

  Forward (primal) value

• tapeIdx : Nat

  Index in the computational tape

def Hesper.AD.Reverse.instInhabitedDual.default : Dual

Equations

• Hesper.AD.Reverse.instInhabitedDual.default = { primal := default, tapeIdx := default }

instance Hesper.AD.Reverse.instInhabitedDual : Inhabited Dual

Equations

• Hesper.AD.Reverse.instInhabitedDual = { default := Hesper.AD.Reverse.instInhabitedDual.default }

instance Hesper.AD.Reverse.instReprDual : Repr Dual

Equations

• Hesper.AD.Reverse.instReprDual = { reprPrec := Hesper.AD.Reverse.instReprDual.repr }

def Hesper.AD.Reverse.instReprDual.repr : Dual → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.Dual.const (x : Float) : Dual

Create a constant (no gradient)

Equations

• Hesper.AD.Reverse.Dual.const x = { primal := x, tapeIdx := 0 }

def Hesper.AD.Reverse.Dual.var (x : Float) (tape : Tape) : Tape × Dual

Create a variable (will have gradient)

Equations

• Hesper.AD.Reverse.Dual.var x tape = match tape.add [] [] with | (tape', idx) => (tape', { primal := x, tapeIdx := idx })

def Hesper.AD.Reverse.Dual.lift1 (f df : Float → Float) (x : Dual) (tape : Tape) : Tape × Dual

Lift a unary function to Dual numbers

Equations

• Hesper.AD.Reverse.Dual.lift1 f df x tape = match tape.add [x.tapeIdx] [df x.primal] with | (tape', idx) => (tape', { primal := f x.primal, tapeIdx := idx })

def Hesper.AD.Reverse.Dual.lift2 (f df_dx df_dy : Float → Float → Float) (x y : Dual) (tape : Tape) : Tape × Dual

Lift a binary function to Dual numbers

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.Dual.lift {Op I O : Type} [inst : Core.Differentiable Op I O] (op : Op) (x : I) (tape : Tape) (parentIndices : List Nat) (localGrads : List Float) : Tape × O × Nat

Lift a generic Differentiable operation to Dual numbers

Equations

• Hesper.AD.Reverse.Dual.lift op x tape parentIndices localGrads = match tape.add parentIndices localGrads with | (tape', idx) => (tape', Hesper.Core.forward op x, idx)

structure Hesper.AD.Reverse.ADContext : Type

AD context for managing the computational tape

• tape : Tape

  The computational tape

def Hesper.AD.Reverse.instInhabitedADContext.default : ADContext

Equations

• Hesper.AD.Reverse.instInhabitedADContext.default = { tape := default }

instance Hesper.AD.Reverse.instInhabitedADContext : Inhabited ADContext

Equations

• Hesper.AD.Reverse.instInhabitedADContext = { default := Hesper.AD.Reverse.instInhabitedADContext.default }

def Hesper.AD.Reverse.ADContext.new : ADContext

Create a new AD context

Equations

• Hesper.AD.Reverse.ADContext.new = { tape := Hesper.AD.Reverse.Tape.empty }

def Hesper.AD.Reverse.ADContext.var (ctx : ADContext) (x : Float) : ADContext × Dual

Create a variable in this context

Equations

• ctx.var x = match Hesper.AD.Reverse.Dual.var x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

def Hesper.AD.Reverse.ADContext.add (ctx : ADContext) (x y : Dual) : ADContext × Dual

Addition

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.ADContext.sub (ctx : ADContext) (x y : Dual) : ADContext × Dual

Subtraction

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.ADContext.mul (ctx : ADContext) (x y : Dual) : ADContext × Dual

Multiplication

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.ADContext.div (ctx : ADContext) (x y : Dual) : ADContext × Dual

Division

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.ADContext.exp (ctx : ADContext) (x : Dual) : ADContext × Dual

Exponentiation

Equations

• ctx.exp x = match Hesper.AD.Reverse.Dual.lift1 Float.exp Float.exp x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

def Hesper.AD.Reverse.ADContext.log (ctx : ADContext) (x : Dual) : ADContext × Dual

Natural logarithm

Equations

• ctx.log x = match Hesper.AD.Reverse.Dual.lift1 Float.log (fun (x : Float) => 1.0 / x) x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

def Hesper.AD.Reverse.ADContext.sqrt (ctx : ADContext) (x : Dual) : ADContext × Dual

Square root

Equations

• ctx.sqrt x = match Hesper.AD.Reverse.Dual.lift1 Float.sqrt (fun (x : Float) => 0.5 / x.sqrt) x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

def Hesper.AD.Reverse.ADContext.sin (ctx : ADContext) (x : Dual) : ADContext × Dual

Sine

Equations

• ctx.sin x = match Hesper.AD.Reverse.Dual.lift1 Float.sin Float.cos x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

def Hesper.AD.Reverse.ADContext.cos (ctx : ADContext) (x : Dual) : ADContext × Dual

Cosine

Equations

• ctx.cos x = match Hesper.AD.Reverse.Dual.lift1 Float.cos (fun (x : Float) => -x.sin) x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

def Hesper.AD.Reverse.ADContext.pow (ctx : ADContext) (x : Dual) (n : Float) : ADContext × Dual

Power (x^n)

Equations

• ctx.pow x n = match Hesper.AD.Reverse.Dual.lift1 (fun (x : Float) => x.pow n) (fun (x : Float) => n * x.pow (n - 1.0)) x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

def Hesper.AD.Reverse.ADContext.sigmoid (ctx : ADContext) (x : Dual) : ADContext × Dual

Sigmoid activation

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.ADContext.tanh (ctx : ADContext) (x : Dual) : ADContext × Dual

Tanh activation

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.ADContext.relu (ctx : ADContext) (x : Dual) : ADContext × Dual

ReLU activation

Equations

• ctx.relu x = match Hesper.AD.Reverse.Dual.lift1 (fun (x : Float) => max 0.0 x) (fun (x : Float) => if x > 0.0 then 1.0 else 0.0) x ctx.tape with | (tape', d) => ({ tape := tape' }, d)

Verified op integration

The methods above (add, mul, sub, …) hardcode each op's local gradient in the body, which is fine for the built-in arithmetic but forces every new op to be hand-coded. liftBinaryVerified plugs the gap by reading the local gradient straight out of the Differentiable instance — i.e. the same forward/backward pair we'd register for any Hesper verified op.

The cost is the function-pointer indirection through the typeclass; the benefit is that any Differentiable Op (Float × Float) Float becomes a Dual op with no further wiring. Used by Ch11 for SquaredErrorOp so the MSE training loop is built from a verified loss op rather than ad-hoc mul/sub/pow.

def Hesper.AD.Reverse.ADContext.liftBinaryVerified {Op : Type} [inst : Core.Differentiable Op (Float × Float) Float] (ctx : ADContext) (op : Op) (x y : Dual) : ADContext × Dual

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.backprop (tape : Tape) (outputIdx : Nat) : Array Float

Reverse-mode gradient computation via backpropagation

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.grad (f : ADContext → Dual → ADContext × Dual) (x : Float) : Float

Compute gradient of a univariate function using reverse-mode AD

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.gradWith (f : ADContext → Dual → ADContext × Dual) (x : Float) : Float × Float

Compute value and gradient of a univariate function

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.gradN (f : ADContext → Array Dual → ADContext × Dual) (xs : Array Float) : Array Float

Compute gradient of a function with respect to multiple variables

Equations

• One or more equations did not get rendered due to their size.

def Hesper.AD.Reverse.gradNWith (f : ADContext → Array Dual → ADContext × Dual) (xs : Array Float) : Float × Array Float

Compute value and gradient for multiple variables

Equations

• One or more equations did not get rendered due to their size.