Differentiable Programming

Differentiable programming treats differentiation as a general programming-language feature. A program can contain numerical kernels, control flow, data structures, solvers, simulations, and model code. If the program denotes a differentiable computation, the system should be able to compute derivatives of that program.

This is broader than neural network training. A differentiable program may include a physics simulator, a renderer, a database query, an optimizer, or a probabilistic model. The central idea is that gradients should be available across the whole program, not only inside a fixed tensor graph.

From Automatic Differentiation to Differentiable Programming

Automatic differentiation began as a technique for computing derivatives of numerical programs. Differentiable programming extends that idea into a design principle for languages and systems.

The core operator is a program transformation:

Given a scalar-output function, grad produces another function that returns its gradient.

In code, this appears as:

def loss(params, batch):
    pred = model(params, batch.x)
    return mse(pred, batch.y)

g = grad(loss)(params, batch)

The important point is not the syntax. The important point is that loss can be an ordinary program. It can call functions, branch, loop, allocate intermediate values, and invoke libraries.

Differentiation as a First-Class Transformation

In a differentiable programming system, transformations such as grad, vjp, jvp, jacfwd, and jacrev are first-class operations.

This style changes how numerical software is written. Instead of manually deriving update rules, the programmer writes the objective and asks the system for derivatives.

Programs Beyond Tensor Graphs

Early deep learning systems often represented models as static tensor graphs. A graph contained operations such as matrix multiplication, convolution, addition, and activation functions.

Differentiable programming relaxes this model. The differentiable object is the program itself.

For example:

def f(x):
    y = 0.0
    for i in range(10):
        if x[i] > 0:
            y = y + x[i] * x[i]
        else:
            y = y - x[i]
    return y

This program contains a loop and a branch. A differentiable programming system differentiates the executed computation. The derivative depends on the path taken through the program.

This makes differentiation useful for real numerical code, where fixed graphs are often too restrictive.

The Language Boundary

A central difficulty is deciding which parts of the language are differentiable.

Arithmetic over real-valued arrays is differentiable. Integer indexing, mutation, allocation, I/O, hashing, sorting, and discrete control decisions need more care.

A differentiable programming language must make these boundaries explicit. Silent behavior is dangerous. A system should reject invalid differentiation, define a subgradient convention, or require a custom rule.

Custom Derivative Rules

Not every useful operation should be differentiated by expanding its implementation. Some operations have better derivative formulas than their source code suggests.

For example, a numerically stable logsumexp implementation may contain branching and shifts:

def logsumexp(x):
    m = max(x)
    return m + log(sum(exp(x - m)))

Differentiating the implementation directly may work, but a custom derivative rule is often clearer and more stable.

A custom rule gives the AD system a local derivative definition:

Custom derivatives are part of the language interface. They allow library authors to expose mathematically correct and efficient differentiation behavior.

Differentiating Through Libraries

A real program uses libraries. A differentiable programming system must decide how derivative information crosses library boundaries.

There are several strategies.

Each strategy has a different boundary. Operator overloading works well when all operations dispatch through overloaded types. Tracing works when the runtime can observe the operations. Source transformation works when source code is available and transformable. Compiler IR differentiation works when the program has been lowered into an analyzable representation.

Differentiable Programming and Compilation

Differentiable programming is closely connected to compiler design.

A compiler-based system can perform several transformations together:

1. Normalize the original program.
2. Differentiate the normalized representation.
3. Optimize the primal and derivative code.
4. Fuse kernels.
5. Plan memory.
6. Lower to CPU, GPU, TPU, or accelerator code.

This is important because naive differentiated programs are often inefficient. Reverse mode can store too many intermediates. Higher-order AD can duplicate computation. Tensor programs can produce many small kernels.

A practical differentiable compiler performs AD and optimization as a single pipeline.

The Role of Types

Types help define which programs are differentiable.

A language may distinguish between:

A typed system can reject invalid uses early:

@differentiable
func f(_ x: Float) -> Float {
    return x * x
}

The annotation marks a function as participating in differentiation. The compiler can then check whether all operations inside the function support derivatives.

Shape types and effect systems extend this further. Shape types ensure tensor dimensions match. Effect systems track mutation, randomness, I/O, and other behavior that affects derivative semantics.

Mutation and State

Mutation is one of the hardest issues in differentiable programming.

Consider:

def f(x):
    y = x
    y[0] = y[0] * 2
    return sum(y)

The assignment changes the value of y. If y aliases x, the mutation also changes x. Reverse mode must reconstruct the correct sequence of states and propagate adjoints through each update.

There are several implementation choices.

Functionalization is common. The system rewrites mutation into pure operations, making the derivative transformation easier to define.

Control Flow and Dynamic Programs

Differentiable programming must support control flow because scientific and machine learning code is full of it.

A loop differentiates the actual iterations run:

def fixed_point(x):
    y = x
    while norm(g(y) - y) > 1e-6:
        y = g(y)
    return y

Differentiating this program by unrolling the loop gives the derivative of the algorithm, not necessarily the derivative of the mathematical fixed point. These can differ.

This distinction matters.

For solvers, implicit differentiation may be preferable to differentiating every iteration.

Differentiable Programming in Practice

Modern systems implement differentiable programming with different tradeoffs.

The field has no single dominant architecture. The right design depends on the host language, execution target, and expected workloads.

Differentiable Programming Beyond Machine Learning

Differentiable programming is useful wherever a program contains parameters that should be optimized.

Examples include:

The common pattern is objective-driven computation. A program computes a loss, error, likelihood, reward, or constraint violation. AD provides the derivative needed to improve parameters.

Design Requirements

A serious differentiable programming system should provide:

Differentiable programming becomes a systems problem, not only a calculus problem.

Failure Modes

Differentiable programming systems fail in characteristic ways.

These failures are often subtle. Good systems expose diagnostics, gradient checks, and explicit boundaries between differentiable and non-differentiable code.

Summary

Differentiable programming generalizes automatic differentiation from numerical kernels to whole programs. It asks the language and compiler to treat derivatives as ordinary program transformations.

The core challenge is semantic clarity. A user should know whether the system differentiates the mathematical function, the algorithm, the implementation, or a custom abstraction exposed by a library. Once that boundary is clear, the remaining problems are compiler and runtime engineering: representation, memory, optimization, dispatch, and hardware lowering.