let f0: @differentiable (Float, Float) -> Float = { $0 * $1 } _ = f0 as @differentiable (Float, @nondiff Float) -> Float _ = f0 as @differentiable (@nondiff Float, Float) -> Float _ = f0 as @differentiable (@nondiff Float, @nondiff Float) -> Float Differential operators The core differentiation APIs are the differential operators. Differential operators are higher-order functions that take @differentiable functions as inputs and return derivative functions or evaluate derivative values. Differential-producing differential operators Among these differential operators, two base APIs, valueWithDifferential(at:in:) and transpose(of:), are used for implementing all other differential operators and differentiation APIs. /// Returns `body(x)` and the differential of `body` at `x`. func valueWithDifferential( at x: T, in body: @differentiable (T) -> R) -> (value: R, differential: @differentiable(linear) (T.TangentVector) -> R.TangentVector) { // Compiler built-in. Builtin.autodiffApply_jvp_arity1(body, x)} /// Returns the transpose of the linear map `body`. func transpose( of body: @escaping @differentiable(linear) (T) -> R) -> @differentiable(linear) (R) -> T { // Compiler built-in. return { x in Builtin.autodiffApply_transpose(body, x) }} The most common differential operators are the ones that compute directional derivatives. These differential operators are defined to take a differentiable function whose parameter is a real number. func valueWithDerivative( at x: T, in body: @differentiable (T) -> R) -> (value: R, derivative: R.TangentVector) where T.TangentVector: FloatingPoint { let (value, df) = valueWithDifferential(at: x, in: body) return (value, df(T.TangentVector(1)))}func derivative( at x: T, in body: @differentiable (T) -> R) -> R.TangentVector where T.TangentVector: FloatingPoint { valueWithDerivative(at: x, in: body).derivative} func derivative( of body: @escaping @differentiable (T) -> R) -> (T) -> R.TangentVector where T.TangentVector: FloatingPoint { return { x in derivative(at: x, in: body) }} Pullback-producing differential operators Unlike directional derivatives, gradients are computed by pullbacks. Based on the differential-producing differential operator valueWithDifferential(at:in:), valueWithPullback(at:in:) is defined as returning the original value and the transpose of the differential, and valueWithGradient(at:in:) is defined as evaluating the pullback at 1 when the function being differentiated returns a real number. func valueWithPullback( at x: T, in body: @differentiable (T) -> R) -> (value: R, pullback: @differentiable(linear) (R.TangentVector) -> T.TangentVector) { let (value, df) = valueWithDifferential(at: x, in: body) return (value, transpose(of: df))} func valueWithGradient( at x: T, in body: @differentiable (T) -> R) -> (value: R, gradient: T.TangentVector) where R.TangentVector: FloatingPoint { let (value, pullback) = valueWithPullback(at: x, in: body) return (value, pullback(R.TangentVector(1)))} func gradient( at x: T, in body: @differentiable (T) -> R) -> (value: R, gradient: T.TangentVector) where R.TangentVector: FloatingPoint { return valueWithGradient(at: x, in: body).gradient} func gradient( of body: @escaping @differentiable (T) -> R) -> (T) -> T.TangentVector where R.TangentVector: FloatingPoint { return { x in gradient(at: x, in: body) }} Example usage All of these APIs are designed to work nicely with Swift’s trailing closure syntax. Here is an example of training a simple deep learning model: for _ in 0..<1000 { // Differentiate the loss with respect to the model `classifier` itself, // producing a tangent vector `𝛁model` that represents partial derivatives // with respect to all differentiable properties (trainable model parameters) // in the model let 𝛁model = gradient(at: classifier) { classifier -> Tensor in let ŷ = classifier(x) let loss = softmaxCrossEntropy(logits: ŷ, labels: y) print("Loss: \(loss)") return loss } optimizer.update(&classifier, along: 𝛁model)} List of differential operators Differential operators Description transpose(of:) Returns transpose of linear map. valueWithDifferential(at:in:) valueWithDifferential(at:_:in:) (arity 2) Returns original result and differential function. valueWithPullback(at:in:) valueWithPullback(at:_:in:) Returns original result and pullback function. differential(at:in:) differential(at:_:in:) (arity 2) Returns differential function. pullback(at:in:) pullback(at:_:in:) Returns pullback function. derivative(at:in:) derivative(at:_:in:) (arity 2) Returns partial derivatives with respect to arguments (“forward-mode”). gradient(at:in:) gradient(at:_:in:) Returns partial derivatives with respect to arguments (“reverse-mode”). valueWithDerivative(at:in:) valueWithDerivative(at:_:in:) (arity 2) Returns original result and partial derivatives with respect to arguments (“forward-mode”). valueWithGradient(at:in:) valueWithGradient(at:_:in:) Returns original result and partial derivatives with respect to arguments (“reverse-mode”). derivative(of:) derivative(of:) (arity 2) Returns derivative function, taking original arguments and returning and partial derivatives with respect to arguments (“forward-mode”). gradient(of:) gradient(of:) Returns gradient function, taking original arguments and returning and partial derivatives with respect to arguments (“reverse-mode”). valueWithDerivative(of:) valueWithDerivative(of:) (arity 2) Returns function taking original arguments and returning original result and partial derivatives with respect to arguments (“forward-mode”). valueWithGradient(of:) valueWithGradient(of:) Returns function taking original arguments and returning original result and partial derivatives with respect to arguments (“reverse-mode”).