Function Transforms

Automatic differentiation transforms for one-argument tensor programs. Write programs against the differentiable operation interface, rather than extracting host values and computing outside the recorded graph. The same program can be interpreted with different scalar types; this is not source-code differentiation of an arbitrary Lean function. Import NN.API.Autograd for the complete public autograd API.

@[reducible, inline]

abbrev TorchLean.autograd.Function (σ τ : Shape) : Type 1

A scalar-polymorphic tensor function written against TorchLean's differentiable operations.

def TorchLean.autograd.Internal.functionProgram {σ τ : Shape} (f : Function σ τ) {α : Type} [Storage α] [Context α] : Runtime.Autograd.Model.Program α [σ] τ

Present a Function as the one-argument Program the autograd runtime consumes.

The wrapping is pure plumbing: curry turns the runtime's heterogeneous argument list into the single ValueRef a Function expects. Every differentiation entry point below goes through here, so exactly one place knows the arity convention.

def TorchLean.autograd.jacfwd {σ τ : Shape} (f : Function σ τ) {α : Type} [Storage α] [Context α] (input : Tensor α σ) : IO (Tensor α (τ.concat σ))

Forward-mode Jacobian with output axes followed by input axes, matching jacrev.

def TorchLean.autograd.hessian {σ : Shape} (f : Function σ []) {α : Type} [Storage α] [Context α] (input : Tensor α σ) : IO (Tensor α (σ.concat σ))

Hessian of a scalar function, with one copy of the input axes for each derivative.

def TorchLean.autograd.vjp {σ τ : Shape} (f : Function σ τ) {α : Type} [Storage α] [Context α] (input : Tensor α σ) (outputGradient : Tensor α τ) : IO (Tensor α σ)

Vector-Jacobian product (VJP) for a pure function.

def TorchLean.autograd.jacrev {σ τ : Shape} (f : Function σ τ) {α : Type} [Storage α] [Context α] (input : Tensor α σ) : IO (Tensor α (τ.concat σ))

Reverse-mode Jacobian (jacrev) of a pure tensor function.

Returns a tensor with output axes followed by input axes, matching jacfwd.

def TorchLean.autograd.grad {σ : Shape} (f : Function σ []) {α : Type} [Storage α] [Context α] (input : Tensor α σ) (value : Bool := false) : IO (match value with | false => Tensor α σ | true => Tensor α σ × Tensor α [])

Differentiate a scalar-valued function with respect to its input.

By default this returns only the gradient. Set value := true to return (gradient, functionValue) from the same evaluation.