autodiff

bun add @stopcock/autodiff

@stopcock/autodiff makes stopcock pipelines differentiable. Each operation records its forward value and local derivative to a scoped tape; gradient() walks the tape backwards and accumulates gradients with the chain rule.

import { pipe } from '@stopcock/fp'
import { differentiable, sin, square, add, type Var } from '@stopcock/autodiff'

const dataFirst = differentiable((x: Var<number>) => sin(add(square(x), 3)))
const dataLast = differentiable((x: Var<number>) => pipe(x, square, add(3), sin))

dataFirst.forward(2) // Math.sin(7)
dataLast.gradient(2) // Math.cos(7) * 4

This first example shows the core idea in both supported call styles. Whether written directly or as a pipeline, the derivative is computed from the same operations, not from a hand-written formula.

Gradient descent playground

Pick a scalar function and watch reverse-mode AD steer each update.

f(x)

What this shows

The green curve is the function. Each orange point is a new guess for x, made by asking autodiff for df/dx and moving in the direction that should lower f(x).

Small learning rates creep toward a valley. Large learning rates can overshoot, bounce, or jump between valleys, which is the same failure mode real optimizers hit.

Function

Start x

Learning rate

Steps

Run Reset

Function Updates

x

f(x)

gradient

next x

Training labs

Two small optimizers driven by live reverse-mode gradients.

Vector fit Contours

Vector fit

The line has two trainable numbers: slope and intercept. Autodiff computes the gradient of mean squared error with respect to both at once, then the demo nudges them until the line fits the sample points.

The lower strip is the loss over time. The outlier preset shows the tradeoff a least-squares model makes when one point fights the trend.

Dataset

Start slope

Start intercept

Learning rate

Steps

Run Reset

Samples Fit Loss

slope

intercept

loss

|gradient|

Contours

This is gradient descent with two variables instead of one. The colored field is the loss surface; the orange path follows the negative gradient in x and y.

Basin is easy, Banana has a narrow curved valley, and Ripple can trap the search in local minima.

Surface

Start x

Start y

Learning rate

Steps

Run Reset

f(x, y)

x

y

loss

|gradient|

Neural classifier

A 4-neuron hidden layer trained live from reverse-mode gradients.

parameters 17

What this shows

The model turns each (x, y) point into a yes/no probability. The colored field is every prediction the model would make across the plane, while the dots are the examples it is trying to classify.

Each training step uses autodiff to measure how all 17 parameters affected the error, then updates them together. XOR and Ring need curved decision boundaries, so a straight line cannot solve them.

Dataset

Seed

Learning rate

Steps

Run Reset

Class 0 Class 1 Loss

loss

accuracy

|gradient|

confidence

Differentiable design tools

Turn a visual goal into parameters by following gradients.

Robot arm Curve fit

Robot arm

This is inverse kinematics. The target is a point in space, and the trainable numbers are the shoulder and elbow angles. Autodiff tells the arm which way to rotate each joint to reduce the wrist-to-target distance.

The dotted trail is the wrist's search path. Move the target or starting angles to see the same objective settle into different valid poses.

Target x

Target y

Start shoulder

Start elbow

Learning rate

Steps

Run Reset

Arm Target Search

shoulder

elbow

distance

|gradient|

Curve fit

The orange dots are target samples. The blue cubic curve has two trainable control points, and autodiff computes how each control coordinate changes the total curve error.

Press Run to watch the handles slide into a shape that best explains the samples. This is the same idea as fitting model weights, but applied to a visual design object.

Target

Learning rate

Steps

Run Reset

Target points Curve Controls

control 1

control 2

loss

|gradient|

Why it works with pipe

pipe is plain, sequential function application. There’s no runtime fusion step to bypass. Each autodiff op runs for real and records to the active tape as it executes. The build-time @stopcock/fp-compiler only recognizes @stopcock/fp/array operators, so a pipe chain of autodiff ops is unaffected either way.

The API

const dataFirst = differentiable((x: Var<number>, y: Var<number>) => add(mul(x, y), square(x)))
const dataLast = differentiable((x: Var<number>, y: Var<number>) =>
  pipe(x, mul(y), add(pipe(x, square))),
)

dataFirst.forward(2, 3) // 10
dataLast.gradient(2, 3) // [7, 2]
dataFirst.valueAndGradient(2, 3) // { value: 10, gradient: [7, 2] }

This example shows how multi-input gradients are returned: one derivative per input, in the same order as the callback parameters.

Callback parameters must be annotated as Var<...>. TypeScript cannot infer the input tuple from inside the callback body, but once the parameter is annotated the returned function is fully typed.

Scalar ops

import {
  add,
  sub,
  mul,
  div,
  neg,
  square,
  pow,
  sin,
  cos,
  tan,
  exp,
  log,
  sqrt,
  abs,
  tanh,
  sigmoid,
  relu,
  leakyRelu,
  softplus,
} from '@stopcock/autodiff'

Operations with two inputs are dual: call them directly with the data first, or partially apply the second input to create a data-last pipe stage.

add(x, 3) // data-first
pipe(x, add(3)) // data-last

leakyRelu(x, 0.2) // data-first
pipe(x, leakyRelu(0.2)) // data-last

Vectors

Vectors are Float64Array values.

import { pipe } from '@stopcock/fp'
import { differentiable, sub, square, vecDot, type Var, type Vec } from '@stopcock/autodiff'

const x = new Float64Array([1, 2, 3])

const dataFirstLoss = differentiable((w: Var<Vec>) => square(sub(vecDot(w, x), 1)))
const dataLastLoss = differentiable((w: Var<Vec>) => pipe(w, vecDot(x), sub(1), square))

dataFirstLoss.gradient(new Float64Array([0, 0, 0]))
dataLastLoss.gradient(new Float64Array([0, 0, 0]))

These equivalent examples show a tiny linear model. vecDot(w, x) and pipe(w, vecDot(x)) both predict a value, square(...) measures the error, and gradient(...) returns a vector shaped like w, telling each weight how to move.

Vector operations include vecAdd, vecSub, vecScale, vecDot, vecNorm, and vecSum.

Matrices

Matrices use the same shape as @stopcock/la:

type Mat = { data: Float64Array; rows: number; cols: number }

import { pipe } from '@stopcock/fp'
import {
  differentiable,
  matMul,
  matNormSquared,
  matSub,
  type Mat,
  type Var,
} from '@stopcock/autodiff'

const target: Mat = { rows: 2, cols: 1, data: new Float64Array([1, -1]) }

const dataFirstLoss = differentiable((w: Var<Mat>, x: Var<Mat>) =>
  matNormSquared(matSub(matMul(x, w), target)),
)
const dataLastLoss = differentiable((w: Var<Mat>, x: Var<Mat>) =>
  pipe(x, matMul(w), matSub(target), matNormSquared),
)

This example shows differentiation through matrix multiplication. The loss is a scalar, but the gradients keep the original matrix shapes, so callers can update model matrices directly.

Matrix operations include matMul, matAdd, matSub, matScale, matTranspose, matSum, matMean, and matNormSquared.

Tape primitives

Use the lower-level API when you need to control the tape scope yourself:

import { withTape, variable, backward, gradOf } from '@stopcock/autodiff'
import { record } from '@stopcock/autodiff/tape'

withTape((tape) => {
  const x = variable(3)
  const parents = [x]
  const doubleBackward = (grad: number) => [grad * 2]

  const y = record(x.value * 2, parents, doubleBackward)
  // Equivalent data-last form:
  // record(parents, doubleBackward)(x.value * 2)

  // Data-first
  backward(y, tape)
  const directGradient = gradOf(x, tape)

  // Equivalent data-last forms
  backward(tape)(y)
  const curriedGradient = gradOf(tape)(x)

  return { directGradient, curriedGradient }
})

This example shows what differentiable(...) wraps for you: create a variable, record operations on a tape, walk the tape backwards, and read the gradient from the original input.

Complete dual reference

The left call is direct and data-first. The right call is the equivalent data-last form.

add(a, b) / add(b)(a)
sub(a, b) / sub(b)(a)
mul(a, b) / mul(b)(a)
div(a, b) / div(b)(a)
pow(a, exponent) / pow(exponent)(a)
leakyRelu(value, alpha) / leakyRelu(alpha)(value)

vecAdd(a, b) / vecAdd(b)(a)
vecSub(a, b) / vecSub(b)(a)
vecScale(vector, scalar) / vecScale(scalar)(vector)
vecDot(a, b) / vecDot(b)(a)

matAdd(a, b) / matAdd(b)(a)
matSub(a, b) / matSub(b)(a)
matMul(a, b) / matMul(b)(a)
matScale(matrix, scalar) / matScale(scalar)(matrix)

accumulate(existing, incoming) / accumulate(incoming)(existing)
record(value, parents, backwardFn) / record(parents, backwardFn)(value)
backward(output, tape) / backward(tape)(output)
gradOf(variable, tape) / gradOf(tape)(variable)

Constraints

• Outputs are scalar in v1.
• Operations are synchronous; do not await inside a differentiable callback.
• Vector and matrix ops throw ShapeError on incompatible shapes.
• Non-differentiable scalar points use the documented subgradient convention: abs(0) = 0, relu(0) = 0, and leakyRelu(0, alpha) = alpha.