AD in Python

Python became the dominant language for modern machine learning and differentiable computing because it combines a simple programming model with access to high-performance native libraries. Most Python automatic differentiation systems therefore follow a hybrid architecture:

Python itself is slow for numerical kernels. The important observation is that tensor operations are executed outside the Python interpreter. The AD system therefore differentiates tensor programs driven by Python control.

Tensor-Centric Computation

Modern Python AD systems are built around tensors.

A tensor object typically contains:

A simple example:

x = tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x * x
z = y.sum()
z.backward()

After execution:

x.grad

contains:

The user writes imperative Python code, but the AD system records tensor operations and constructs derivative propagation rules internally.

Dynamic Computational Graphs

Many Python systems use dynamic graphs.

A graph is built during execution. Each tensor operation creates graph nodes connecting inputs to outputs.

For example:

a = x + 1
b = sin(x)
c = a * b

produces a graph:

x
├── add ── a
├── sin ── b
└── mul(a,b) ── c

Each node stores:

Dynamic graphs are flexible because they naturally support:

• Loops
• Branches
• Recursion
• Variable shapes
• Interactive execution

This matched the needs of machine learning research, where models evolve rapidly.

Reverse Mode in Python Systems

Most Python AD frameworks optimize for scalar-loss reverse mode because neural network training requires gradients with respect to many parameters.

Suppose:

y = f(x)

where:

Reverse mode computes:

with cost proportional to a small multiple of the forward evaluation.

Internally, the reverse pass traverses the graph backward:

1. Initialize output adjoint to 1.
2. Visit graph nodes in reverse topological order.
3. Apply local derivative rules.
4. Accumulate gradients into parent tensors.

For multiplication:

the reverse rule is:



The Python frontend hides this machinery, but the runtime manages graph traversal, tensor lifetimes, and adjoint accumulation.

Eager Execution

Many Python systems use eager execution.

Operations execute immediately:

x = tensor([1, 2, 3])
y = x + 1

y is computed immediately rather than deferred into a static graph.

Advantages include:

The downside is optimization difficulty. Since the runtime sees operations incrementally, it has limited global visibility.

Tracing and Graph Capture

To recover optimization opportunities, many systems trace Python functions into graph representations.

Example:

def f(x):
    return sin(x) + x * x

Tracing executes the function with special tensor objects that record operations instead of performing ordinary computation.

The result is an intermediate graph:

x
├── sin
├── mul(x,x)
└── add

The graph can then be optimized, fused, compiled, or lowered to accelerators.

Tracing enables:

This produces a hybrid execution model:

Static vs Dynamic Graph Systems

Early Python AD systems often used static graphs.

The user first defined a graph:

x = placeholder()
y = x * x

and later executed it:

session.run(y, feed_dict={x: ...})

Static graphs enabled aggressive optimization but created awkward programming models.

Dynamic systems later became dominant because they matched ordinary Python execution.

The distinction today is less strict. Modern systems often combine both:

Mutation and In-Place Operations

Python tensor systems frequently support mutation:

x += y

Mutation complicates reverse mode because the old value of x may be needed during the backward pass.

Systems handle this differently.

PyTorch, for example, tracks tensor versions to detect modifications that invalidate gradient computation.

Custom Gradient Functions

Many operations need manually defined derivatives.

A Python framework usually exposes an API:

class MyOp(Function):

    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return ...

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        return ...

This separates:

Custom gradients are critical for:

• Numerical stability
• Efficient kernels
• External libraries
• Implicit differentiation
• Physics simulators
• Specialized GPU operations

Higher-Order Differentiation

Python systems increasingly support higher-order derivatives.

Example:

g = grad(f)
h = grad(g)

This requires differentiating the backward pass itself.

The system must ensure:

• Reverse-mode operations are differentiable
• Graphs remain valid through nesting
• Saved tensors survive nested passes
• Perturbation confusion is avoided

Higher-order AD is important for:

Python and Functional Transformations

Some Python AD systems adopt a more functional style.

Instead of mutating tensor objects:

y.backward()

they expose explicit transformations:

grad(f)
vmap(f)
jit(f)
pmap(f)

These transformations compose.

Example:

jit(grad(f))

means:

1. Differentiate f
2. Compile the resulting derivative program

This model treats differentiation as a pure program transformation rather than as a side effect attached to tensors.

Compilation Pipelines

Modern Python AD frameworks often lower programs into compiler IRs.

The pipeline may look like:

Python
→ traced graph
→ normalized IR
→ optimized IR
→ backend lowering
→ machine code or accelerator kernels

Common IR forms include:

AD may operate at multiple levels:

Interaction with NumPy

NumPy heavily influenced Python AD systems.

Many frameworks mimic NumPy APIs:

sin
exp
matmul
reshape
broadcast_to
sum
transpose

This allows numerical code to become differentiable with minimal changes.

However, ordinary NumPy arrays do not carry gradient metadata. Frameworks therefore provide tensor types that emulate NumPy behavior while tracking derivatives.

Compatibility layers are essential for ecosystem adoption.

GPU and Accelerator Execution

Python frameworks usually execute tensor kernels outside Python.

The Python interpreter orchestrates computation, but dense operations are dispatched to:

AD systems must therefore manage:

• Device placement
• Gradient synchronization
• Memory transfers
• Kernel scheduling
• Mixed precision

The derivative computation becomes part of a distributed runtime system.

Memory Management

Reverse mode requires storing intermediates from the forward pass.

Memory costs can dominate execution.

Strategies include:

Large neural networks are often constrained more by activation memory than by arithmetic throughput.

Numerical Stability

Naive derivatives can be numerically unstable.

Examples include:

Python AD systems therefore rely heavily on custom stable primitives.

For example:

logsumexp
softplus
cross_entropy
layer_norm

often have carefully engineered backward implementations.

Major Python AD Systems

These systems differ mainly in:

• Graph construction strategy
• Compilation model
• Mutation semantics
• Transformation interface
• Hardware integration

Python as an AD Host Language

Python succeeded because it provided:

The AD engine is usually not written primarily in Python. Python acts as the orchestration layer above highly optimized native runtimes and compiler systems.

Modern Python AD frameworks therefore resemble compiler toolchains hidden behind an imperative scripting interface.