Beyond Backpropagation - Higher Order, Forward and Reverse-mode Automatic Differentiation for Tensorken
Tags: rust, autodiff, tensors, ad
Published: 2023-12-10 12:02:05

---
title: Beyond Backpropagation - Higher Order, Forward and Reverse-mode Automatic Differentiation for Tensorken
published: true
date: 2023-12-09 22:00:04 UTC
tags: rustlang,autodiff,tensors,ad
canonical_url: https://getcode.substack.com/p/beyond-backpropagation-higher-order
---

This post describes how I added [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation) to [Tensorken](https://github.com/kurtschelfthout/tensorken). Tensorken is my attempt to build a fully featured yet easy-to-understand and hackable implementation of a deep learning library in Rust. It takes inspiration from the likes of [PyTorch](https://pytorch.org/), [Tinygrad](https://github.com/tinygrad/tinygrad), and [JAX](https://jax.readthedocs.io/en/latest/).

Tensorken's approach to automatic differentiation (or AD) is heavily inspired by JAX. Like JAX, Tensorken supports higher-order derivatives - besides the first derivative, it can calculate the second, third, and so on. Tensorken supports both forward and reverse-mode AD, and can arbitrarily compose the two. Finally, thanks to good fundamentals explained in the previous two posts ([part 1](https://getcode.substack.com/p/fun-and-hackable-tensors-in-rust) and [part 2](https://getcode.substack.com/p/massively-parallel-fun-with-gpus)), Tensorken can compute derivatives on the CPU or GPU.

All code for this post is in the [Tensorken](https://github.com/kurtschelfthout/tensorken) repository, [tagged v0.3](https://github.com/kurtschelfthout/tensorken/tree/v0.3).

![Four turtles on top of each other on a forest background](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/3me3r78vq8itlx16tsgx.jpg)
_There's a "it's turtles all the way down" reference somewhere in this post, and only then will this image make sense. Generated by Lexica's Aperture model._

## **Previously in Tensors From Scratch: neural networks, matrix multiplication, and GPU acceleration**

Modern neural networks, for example, large language models (LLMs) like OpenAI's ChatGPT and GPT-4, Microsoft's Bing, Google's Bard, and Anthropic's Claude, are powered by tensors. Tensors are multi-dimensional arrays augmented with operations that execute efficiently on modern hardware, most notably GPUs.

To understand all that, I am building a neural network library like PyTorch or JAX, from the ground up in Rust. These libraries consist of:

1. A tensor library, to provide efficient operations to slice, dice, and apply bulk operations to tensors.

2. Accelerators, to accelerate tensor operations on the GPU.

3. Automatic differentiation, to train neural networks via gradient descent.

4. Neural network building blocks, to simplify using common activation functions and layers.

In [the first post](https://getcode.substack.com/p/fun-and-hackable-tensors-in-rust), I focussed on the tensor library. I described almost twenty fundamental tensor operations and abstracted them in a Rust trait called `RawTensor`. `RawTensor` had a single implementation, `CpuRawTensor` which executes tensor computations on the CPU. In [the second post](https://getcode.substack.com/p/massively-parallel-fun-with-gpus), I implemented `RawTensor` again in `WgpuRawTensor` to execute on the GPU using [wgpu](https://wgpu.rs/web), Rust's implementation of [WebGPU](https://www.w3.org/TR/webgpu/). We dove into the nitty-gritty of GPU programming in general and wgpu in particular.

This third part of the series describes how to add automatic differentiation to Tensorken. Automatic differentiation (AD) is a technique to compute derivatives of tensor computations, without programmer intervention. AD is crucial because neural networks are trained via gradient descent, which relies on the efficient calculation of derivatives.

## How to train your neural network

Let's sketch how to train a neural network to emphasize how important AD is for deep learning.

First, gather training data, and lots of it. Training data are lots of input-expected output pairs. The input examples are encoded as numbers and aggregated in a tensor 𝚇. The outputs go in a tensor 𝚈. Think of 𝚈 as the correct predictions for the inputs 𝚇. For a language model, each example in 𝚇 could be a sequence of words, and 𝚈 the next word, encoded as numbers. (How to encode text as numbers is an interesting problem that's not relevant to this story.)

Second, decide on the architecture of your neural network. A neural network consists of tensors 𝚆ᵢ that contain the parameters of the network. That's what you download when you get a model's weights. The architecture determines how many parameters we have and how we combine the input 𝚇 with parameters Wᵢ to obtain an output 𝚈'. Whatever the architecture is, we can execute it to predict 𝚈':

𝚈' = 𝚏(𝚇, 𝚆ᵢ).

I'm simplifying - researchers distinguish weights 𝚆 and biases b, but in the end, they're both part of the trainable parameters so I'm just lumping them together in the 𝚆s.

Third, using the expected 𝚈 and the prediction 𝚈', calculate the loss 𝙻. The loss is a single number that is high when the prediction is bad, and low when it is good. The loss is calculated by comparing the network's prediction 𝚈' with the expected output 𝚈:

𝙻 = 𝚕(𝚈, 𝚈').

Fourth, calculate the gradient 𝙶 of the loss. The scalar loss value 𝙻 is a function of 𝚇, 𝚆ᵢ, and 𝚈. Imagine the loss function as describing a (highly dimensional!) landscape. Training the network to improve its predictions means changing the parameters Wᵢ to make the loss small. We'd like to know _how_ we should change the parameters to achieve that.

Now is when the gradient comes in. Going back to the landscape analogy, to make the loss smaller we'd like to know the best direction to "move" in to go "down" - that is, from the current value of the parameters, find the direction with the highest slope. If you remember some calculus, the derivative of a function at a point is that slope. So, to calculate the gradient, we calculate the loss function's derivative with respect to each parameter 𝚆ᵢ. In other words, we'll have a number for each parameter that tells us how to change that parameter to make the loss smaller.

Fifth, update the parameters using the gradient. There are many ways of doing this. The simplest is to multiply the gradient with a small number ϵ and subtract it from the parameters:

Wᵢ <- Wᵢ − ϵ𝙶ᵢ.

That's one training step done! Your neural network just got a tiny bit better. Now repeat from step 2 until you've had enough. You can stop when the loss becomes small enough, when it stops changing for some number of iterations, or when your AWS bill exceeds the budget.

Tensorken can already do almost all of those steps. Running a network, calculating a loss, and updating the parameters amounts to applying tensor operations. What's missing is calculating the gradient via the loss function's derivative. In the olden days, people would calculate the derivative of the network by hand, symbolically, and then implement it manually. Clearly tedious and error-prone, not to mention limiting the complexity and size of the networks. Modern neural network libraries calculate a function's output and its derivative without programmer intervention using the miracle of automatic differentiation.

AD is a vast and intricate topic. For a (much) longer primer on the basics, see my [earlier post](https://getcode.substack.com/p/automatic-differentiation-from-forward). If you are unfamiliar with AD I encourage you to read it or any of the AD primers in the links.

The following section demonstrates Tensorken's AD capabilities and interface via small examples. Then we'll dive into implementation details, but I'll stay away from the detailed mechanics of AD since that's already covered elsewhere. Instead, I'll focus on how Tensorken implements higher-order, mixed-mode, JAX-style AD as an elegant and minimal Rust library.

## The Autodiff Cookbook in Tensorken

To demonstrate Tensorken's AD capabilities, I translated a significant part of JAX's [Autodiff Cookbook](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html) to Tensorken. I reproduced and edited part of the original text here. JAX's license is Apache 2.0, so I hope this does not incur the wrath of Google. The titles in this section are similar to the ones in JAX's cookbook, in case you want to compare. The full example code is in [jax\_autodiff\_cookbook.rs](https://github.com/kurtschelfthout/tensorken/blob/v0.3/examples/jax_autodiff_cookbook.rs).

Before we begin - Tensorken runs on the CPU if you create tensors via the `Cpu32` type alias and on the GPU via `Wgpu32`. (The `32` is because they work with 32-bit floating point numbers.) To make it easy to switch, I'll use the `Tr` type alias throughout:

```rust
type Tr = Cpu32; // or Wgpu32
```

### Gradients

You can differentiate a function using `grad1`. The `1` indicates the number of arguments of the function - a poor man's variadic arguments. In the text, I'll sometimes refer to the family of `grad1`, `grad2` functions as `grad`. In the code, I'll use the function with the correct number of arguments.

To start with, we'll use a simple scalar function - a function that takes a single number and returns a single number:

```rust
let p = Tr::scalar(2.0);
let df = grad1(|x| x.tanh(), &p);

> df: [0.07065082]
```

In Tensorken, all arguments must be a tensor `Tr` - it doesn't support mixed tensors and scalar numbers. To turn a number into a tensor we first use `Tr::scalar`. It makes a tensor with shape `[1]`.

`grad1` takes a function of one argument 𝚏 and evaluates ∇𝚏(𝚙), the derivative of 𝚏 at a given point 𝚙. You can think of ∇ as a higher-order function that takes a differentiable function 𝚏 and produces a function that evaluates the derivative.

_Pronouncing ∇: I say "grad", I've heard people say "del", and the symbol's Unicode name is "nabla"._

Similarly, if you have a Rust function `f` that evaluates the mathematical function 𝚏, then `grad(f, p)` computes the value ∇𝚏(𝚙).

Unlike JAX, Tensorken does not directly expose ∇𝚏 as a first-class function, mostly because I had a hard time accomplishing that in Rust and staying sane! It required returning a closure from `grad(f)` so you can write `grad(f)(p)`, but satisfying the compiler proved difficult. So far this hasn't been a constraint in practice.

Like JAX, Tensorken does support applying `grad` to functions that themselves call `grad` to calculate higher-order derivatives:

```rust
let ddf = grad1(|x| grad1(|x| x.tanh(), x), &p);
let dddf = grad1(|x| grad1(|x| grad1(|x| x.tanh(), x), x), &p);

> ddf: [-0.13621868]
> dddf: [0.25265408]
```

Let’s try computing gradients with `grad` in a linear logistic regression model. In other words, a simple neural network with one neuron. First, the setup:

```rust
// Outputs probability of a label being true.
fn predict<'t, T>(w: &'t T, b: &'t T, inputs: &T) -> T
where
    T: TensorLike<'t>,
{
    (inputs.dot(w) + b).sigmoid()
}
```

The function `predict` encodes the architecture of our toy model. Its parameters are a vector `w` and a scalar `b`, for weights and bias. As you can see, we're multiplying the weights with the inputs and adding the bias. Then we use `sigmoid` to squish the output values in the [0, 1] interval. This model predicts the probability of an outcome based on some input measurements.

Why is this equivalent to a neural network with a single neuron? Say the vector `w` has three elements - three weights. We thus have three inputs as well, in `inputs`. The function `dot` multiplies each input with its corresponding weight, and then adds them up. The bias `b` in the neuron analogy is typically a negative number, which represents a threshold that `inputs.dot(w)` must exceed to "activate" the neuron.

All arguments are tensors, but are represented by a generic argument `T`. The type needs to be generic so automatic differentiation can work. We'll see later why. `T: TensorLike` is a handy constraint to make tensor operations like `dot`, `+`, and `sigmoid` available on `T`. You'll see the `TensorLike` constraint often when using Tensorken's AD: to make functions differentiable, replace concrete `Tensor` types with `T: TensorLike`.

Let's run the model.

```rust
// Build a toy dataset.
// These are four measurements of some unspecified variable, one in each row.
let inputs = Tr::new(
    &[4, 3],
    &[
        0.52, 1.12, 0.77, //
        0.88, -1.08, 0.15, //
        0.52, 0.06, -1.30, //
        0.74, -2.49, 1.39,
    ],
);
// These are four observed outcomes, one for each row in the input.
let targets = Tr::new(&[4], &[1.0, 1.0, 0.0, 1.0]);

// Initialize the parameters w and b randomly
let key = 0;
let mut rng = StdRng::seed_from_u64(key);
let w = Tr::randn(&[3], &mut rng);
let b = Tr::randn(&[1], &mut rng);

let prediction = predict(&w, &b, &inputs);

> prediction: [0.4059896 0.37711427 0.9770815 0.007901279]
```

The inputs could be "changes in temperature observed on three consecutive days" and the targets could be "temperature went up or down on the next day". We're then training a model that predicts the probability of the temperature going up given three days' changes in temperature.

Since we initialized the model randomly, its prediction is random. We got unlucky: if you compare `targets` (what we want) with `prediction` (what we have) there is a big difference. The 3rd and 4th predictions are especially bad, almost the exact opposite of the training data.

To improve our model, we first need to quantify how crap the model is via a loss function.

```rust
// Training loss is the negative log-likelihood of the training examples.
fn loss<'t, T>(w: &'t T, b: &'t T, inputs: &T, targets: &'t T) -> T
where
    T: TensorLike<'t>,
    for<'s> &'s T: TensorLikeRef<T>,
{
    let prediction = predict(w, b, inputs);
    // ones_like makes a tensor of the same shape with all values equal to 1.
    let label_probs = &prediction * targets
        + (&prediction.ones_like() - &prediction) * (targets.ones_like() - targets);
    -label_probs.log().sum(&[0])
}

let l = loss(&w, &b, &inputs, &targets);

> loss: [10.4931755]
```

This loss function is _negative log-likelihood_. You can intuit why it works: `prediction` is "compared" with `targets` in `label_probs`. It contains a high value for predictions that are close to the target. We then take the `log` of each, which exaggerates its value: the logarithm is -infinity when `label_probs` is zero. Since the logarithm is negative, we negate it to get a positive number. Then we take the `sum` of the vector so we have a single positive loss number that is high when the model is doing badly, and low when it's making good predictions.

Now we can improve the model by adjusting its weights and biases. We use `grad` to differentiate the `loss` function with respect to the parameters `w` and `b`:

```rust
// Differentiate loss wrt weights
let w_grad = grad1(
    |w| {
        loss(
            w,
            &Reverse::lift(&b),
            &Reverse::lift(&inputs),
            &Reverse::lift(&targets),
        )
    },
    &w,
);
print!("w_grad: {w_grad}");

// Differentiate loss wrt bias
let b_grad = grad1(
    |b| {
        loss(
            &Reverse::lift(&w),
            b,
            &Reverse::lift(&inputs),
            &Reverse::lift(&targets),
        )
    },
    &b,
);

> w_grad: [-1.0830948 2.5363755 -3.2000453]
> b_grad: [-1.2319121]
```

To make the types work out, we need to `Reverse::lift` all the arguments to `loss` we do NOT want to differentiate. They are treated as constants. The type is called `Reverse` because Tensorken uses reverse mode AD in this case. The `Reverse` type reveals why we need to make the arguments to `loss` and `predict` generic: the `grad` function, while taking a plain `Tr` type as the second argument, passes `Reverse<Tr>` to the closure. So the function `f` can be called with `Tr`, `Reverse<Tr>`, or other types we'll see later.

Here's the simplified signature for `grad1`.We'll get to the full signature later:

```rust
pub fn grad1<F>(f: F, at: &Tr) -> Tr where F: Fn(&Reverse<Tr>) -> Reverse<Tr>
```

Briefly, `Reverse` is a wrapper to interpret tensor operations so they calculate the derivative along with the main result. In this example, it'll run a different `dot`, `+`, and `sigmoid` compared to calling `loss` with plain tensors of type `Tr`.

Calling the `loss` function twice is not ideal - we're doing twice the work. We can also calculate the gradients with respect to both `w` and `b` at the same time, using `grad2`.

```rust
let (w_grad, b_grad) = grad2(
    |w, b| loss(w, b, &Reverse::lift(&inputs), &Reverse::lift(&targets)),
    &w,
    &b,
);

> w_grad: [-1.0830948 2.5363755 -3.2000453]
> b_grad: [-1.2319121]
```

Finally, let's do a single training iteration and check if that improves our model.

```rust
// Update parameters
let new_w = &w - &w_grad;
let new_b = &b - &b_grad;

// Predict
let new_prediction = predict(&new_w, &new_b, &inputs);
let new_loss = loss(&new_w, &new_b, &inputs, &targets);

> new_prediction: [0.7384342 0.99262685 0.7747804 0.9996524]
> new_loss: [1.8016509]
```

A massive improvement - we're now only 1.8 crap, down from 10.5!

### Evaluate a function and its gradient using `value_and_grad`

In a real training run, we'd do the above in a loop while keeping an eye on the loss to see when to stop. Again `loss` is called twice: once inside `grad` and once outside. Luckily, we don't have to. Another convenient family of functions is `value_and_grad` to efficiently compute a function and its gradient.

```rust
let (loss_value, (w_grad, b_grad)) = value_and_grad2(
    |w, b| loss(w, b, &Reverse::lift(&inputs), &Reverse::lift(&targets)),
    &w,
    &b,
);

> loss: [10.4931755]
> w_grad: [-1.0830948 2.5363755 -3.2000453]
> b_grad: [-1.2319121]
```

### Checking against numerical differences

Our loss improved, which is a good indication that things work. To gain confidence we can compare Tensorken's derivatives with finite differences.

```rust
// step size for finite difference
let eps = Tr::scalar(1e-4);
let half_eps = &eps / Tr::scalar(2.);
let b_grad_numerical = (loss(&w, &(&b + &half_eps), &inputs, &targets)
    - loss(&w, &(&b - &half_eps), &inputs, &targets))
    / &eps;

> b_grad_numerical [-1.2207031]
> b_grad_autodiff [-1.2319121]
```

Close enough.

### Jacobians using `jacfwd` and `jacrev`

Ignoring bias `b` for now, the `loss` function is a function of three parameters, represented as a single tensor `w` with three elements. It has a single scalar output, represented as a tensor with a single element. Taking the gradient of this function results in a vector of three elements, the sensitivity of the loss to each parameter. This picture becomes more complicated if there is more than one output parameter. `grad` still gives an answer, but what does it mean?

```rust
let deriv = grad1(
    |w| predict(w, &Reverse::lift(&b), &Reverse::lift(&inputs)),
    &w,
);

> deriv: [0.34956074 -0.0017646346 0.20271438]
```

Remember that `predict` returns a vector with four elements, and the input `w` is a vector with three elements. We get a vector with three sensitivities - one for each input. But the sensitivity of which output? There are four. As we'll check below, `grad` returns the sum of the sensitivity of all outputs. That's typically not what we want: we'd like to disaggregate the sensitivities.

The usual approach is to represent the sensitivity of each output with respect to each input as a matrix, called the [Jacobian](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant). In this case, a 4 by 3 matrix - number of outputs by number of inputs. Tensorken can compute Jacobians, in forward and reverse mode using `jacfwd` and `jacrev`:

```rust
let J = jacfwd(
    |w| predict(w, &Forward::lift(&b), &Forward::lift(&inputs)),
    &w,
);

> jacfwd result, with shape [4, 3]
┌ ┐
│ 0.12540425 0.2701015 0.18569478 │
│ 0.20671119 -0.25369102 0.03523486 │
│ 0.01164451 0.0013435973 -0.029111274 │
│ 0.0058007482 -0.019518733 0.010895999 │
└ ┘

let J = jacrev(
    |w| predict(w, &Reverse::lift(&b), &Reverse::lift(&inputs)),
    &w,
);

> jacrev result, with shape [4, 3]
┌ ┐
│ 0.12540427 0.27010152 0.18569478 │
│ 0.20671119 -0.25369102 0.03523486 │
│ 0.01164451 0.0013435973 -0.029111274 │
│ 0.005800748 -0.019518731 0.010895998 │
└ ┘
```

These two functions compute the same values (up to machine precision), but differ in their implementation: `jacfwd` uses forward-mode automatic differentiation, which is more efficient for "tall" Jacobian matrices, while `jacrev` uses reverse-mode, which is more efficient for "wide" Jacobian matrices. For matrices that are near-square, `jacfwd` probably has an edge over `jacrev`.

We can now check that `grad` computed the sum of the sensitivity of all four outputs:

```rust
&J.sum(&[0])

> [0.34956074 -0.0017646346 0.20271438]
```

Using a composition of `javfwd` and `jacrev` gives us a way to compute dense Hessian matrices. Hessian matrices contain all the second derivatives.

```rust
let hessian = jacfwd(
    |w| {
        jacrev(
            |w| {
                predict(
                    w,
                    &Reverse::lift(&Forward::lift(&b)),
                    &Reverse::lift(&Forward::lift(&inputs)),
                )
            },
            w,
        )
    },
    &w,
);
println!("hessian with shape {:?}", hessian.shape());

> hessian shape [4, 3, 3]
```

Why this shape? We start with a function f:𝙽→𝙼. Traditionally, we'd write 𝚏:ℝⁿ→ℝᵐ, but that there are 𝙽 inputs and 𝙼 outputs is more important than that we're talking about real numbers, so I'll omit the ℝ from now on.

At a point 𝚡 ∈ 𝙽 we expect to get the shapes

- 𝚏(𝚡) ∈ 𝙼, the value of 𝚏 at 𝚡,

- ∂𝚏(𝚡) ∈ 𝙼 × 𝙽, the Jacobian matrix at 𝚡,

- ∂²𝚏(𝚡) ∈ 𝙼 × 𝙽 × 𝙽, the Hessian at 𝚡,

and so on.

To implement `hessian` we could have used `jacfwd(jacrev(f))` or `jacrev(jacfwd(f))` or any other composition of the two. But forward-over-reverse is typically the most efficient. That’s because, in the inner Jacobian computation, we’re often differentiating a function with a wide Jacobian (maybe like a loss function 𝚏:𝙽→1), while in the outer Jacobian computation, we’re differentiating a function with a square Jacobian (since ∇𝚏:𝙽×𝙽), which is where forward-mode wins out.

Note we now need to `lift` the inputs twice to make Rust's type checker happy.

That concludes the tour of Tensorken's AD capabilities. It packs a lot of punch - now let's see how to fit it in a small package.

## A tale of two functions

All AD functions like `jacfwd`, `jacrev`, and `grad` are implemented in terms of two function-type pairs: `jvp` with `Forward`, and `vjp` with `Reverse`. JVP stands for Jacobian-vector product, and VJP stands for Vector-Jacobian product. These functions are directly inspired by JAX. To explain their names, we need some math background that deserves a standalone post. If you can't wait, refer to [this section](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#how-it-s-made-two-foundational-autodiff-functions) in JAX's Autodiff Cookbook.

I'll now introduce `jvp` and `vjp`, and the beginnings of how AD works in Tensorken. I assume some background knowledge about AD, in particular AD for scalar functions. See [my earlier post](https://getcode.substack.com/p/automatic-differentiation-from-forward) for a primer.

### From scalars to tensors

Forward AD on scalar functions works by replacing operators and functions on numbers with versions that operate on a _dual number_ - a `(f32, f32)` tuple. The first element is the _primal_, which the function computes without AD. The second is the derivative, or _tangent_. Operations on dual numbers are straightforward:

- apply the operation to the primal(s), and

- apply differentiation rules to the tangent(s).

For example, multiplication on dual numbers is:

(p₁, t₁) . (p₂, t₂) = (p₁.p₂, p₁.t₂ + p₂.t₁)

Reverse mode is more involved. The primal computation is identical, but instead of calculating the tangent alongside the primal, we collect a trace - essentially a stack of operations. A reverse pass through the trace calculates the derivatives.

Exactly how these operations are replaced is a concern for the implementation. Common methods are code transformation in the compiler, code generation, and operator overloading. Tensorken uses trait-based overloading.

Forward mode has little extra memory requirements beyond bringing the tangent along for the ride, while reverse mode needs to keep a trace that's as long as the computation is deep. As a result, for scalar-to-scalar functions, forward mode is more efficient.

That situation changes if we consider functions from many scalars to one, or vice versa. One extreme is a function that takes a single input and computes n outputs. That's great for forward mode: in one execution of the function on dual numbers, we'll have both the primal result and the derivative - or in other words the sensitivity of each output to a small change in the single input.

However, a function that takes many inputs and has a single output is efficient only in reverse mode. In forward mode, we'd need as many executions of the function as there are inputs - we'd have to pass 1 as tangent for each input separately. In reverse mode, we still need the extra memory for the trace, but one forward pass for the primal and one backward pass for the partial derivatives is all we need.

The good news is that if you understand this, nothing much changes if we allow tensors instead of scalars. After all, a tensor is a container of scalars, and operations on tensors can be broken down into operations on scalars. That's not how we want to implement them though! Bulk operations are where the performance is at.

For forward mode, we'll overload tensor operations to propagate a "dual tensor", a tuple of a primal and a tangent tensor. For reverse mode, we'll build up a trace of tensor operations in the forward pass and get the tangent tensors from a backward pass.

One difference with scalar AD is that we need to take the shape of tensors into account. Besides arithmetic operations like addition and multiplication, we also need to figure out differentiation rules for `sum`, `reshape`, and others, which affect the shape of both primal and derivative tensors.

### JVP for forward-mode AD

Here's the signature of `jvp`:

```rust
pub fn jvp1<T: Diffable + Clone, F>(f: F, at: &T, tangent: &T) -> (T, T)
where
    for<'a> F: Fn(&'a Forward<T>) -> Forward<T>,
```

As the `1` suffix indicates, this version takes a single primal tensor `at` and a single tangent tensor `tangent`. It evaluates the primal and tangent of the function `f` and returns them as a tuple.

To understand why this signature makes sense, think of AD as a program transformation. Without AD we'd write programs that boil down to:

```rust
let p1 = f1(x);
let p2 = f2(p1);
...
```

With forward AD and `jvp` we can rewrite them as:

```rust
let (p1,t1) = jvp(f1, x, x.ones_like());
let (p2,t2) = jvp(f2, p1, t1);
...
```

That illustrates how programs that compose functions can be transformed into programs that compose `jvp`-wrapped functions.

Importantly, `at` and `tangent` must have the same shape, and the two tensors in the output tuple have the same shape. `f` computes `out.0` from `at`, and `jvp` additionally computes `out.1` from `at` and `tangent`.

`jvp` works for any tensor-like type that implements `Diffable`. `Diffable` is the foundational trait that defines Tensorken's primitive, differentiable tensor operations. Higher-level operations like `matmul` are built on these. Keeping the tensor generic in `jvp` allows it to work with any `Diffable` implementation - something we'll use when doing higher-order AD. We'll come back to the details soon.

### VJP for backward AD

Here is the signature for `vjp`:

```rust
pub fn vjp1<'b, 't, T: Diffable + Clone + 't, F>(f: F, at: &T) -> (T, PullBack<'t, T>)
where
    for<'a> F: Fn(&'a Reverse<'a, 't, T>) -> Reverse<'a, 't, T>,
```

It looks a bit different because reverse mode has a backward pass. What's the same are the differentiable function `f: F` and the primal input `at`. `vjp` calls `f` with `Reverse` wrapping the input `T`. Since reverse mode needs two passes, `vjp` only returns the primal directly.

`PullBack` (a term from differential geometry) is a named struct that executes the backward pass. It takes a _cotangent_, a tensor in the shape of the _output_ of `f`, and calculates the tangent, a tensor in the shape of the input of `f`.

```rust
impl<T: Diffable + Clone> PullBack<'_, T> {
    pub fn call(&self, cotangent: &T) -> T
}
```

It's all backward! But that's why reverse mode AD is more efficient if you have the right tensor shape.

A short note on why `jvp` and `vjp` have different signatures. On the one hand, we could re-write `jvp` to return a `PushForward` struct with a `call` function that works similarly to `jvp`'s `PullBack`. However, that would require keeping a trace of the operations around so users can `call` multiple times with different tangents. That jeopardizes the memory efficiency of forward mode. The ability to re-execute the differentiating pass with different tangents does not offset the added memory usage.

We could also write `vjp` with a signature like `jvp` by making the `PullBack` internal and `call`ing at the end. In reverse mode, we have to expend the memory anyway, so we might as well make it available to the user for potential reuse.

## Interpreters all the way down

We're now at the point where we can dive into the code, and it's interpreters all the way down.

Before AD, Tensorken's core was the `RawTensor` trait, with implementations for the CPU and the GPU. It's useful to think of this trait as the definition of a language for primitive tensor operations, and implementations of the trait as interpreters of that language. Interpreters don't necessarily have to produce a tensor - for debugging and testing a pretty-printing interpreter for `RawTensor` is useful:

```rust
impl RawTensor for String {
    type Elem = f32;

    fn exp(&self) -> Self {
        format!("{self}.exp()")
    }

    fn add(&self, other: &Self) -> Self {
        format!("({self} + {other})")
    }

    // etc
}
```

We can use it as follows:

```rust
let t1: String = RawTensor::new(&[2, 2], &[1., 2., 3., 4.]);
let t2: String = RawTensor::new(&[2, 2], &[5., 6., 7., 8.]);
let r = t1.exp().add(&t2.log());

> r: "(new([2, 2], [1.0, 2.0, 3.0, 4.0]).exp() + new([2, 2], [5.0, 6.0, 7.0, 8.0]).log())"
```

Or even:

```rust
let t1: String = "A".to_string();
let t2: String = "B".to_string();
let r = t1.exp().add(&t2.log());

> r: "(A.exp() + B.log())"
```

We could generate source code or an abstract syntax tree this way, turning the interpreter into a compiler of sorts. That is the essence of the final tagless approach I described in depth [in an earlier post](https://getcode.substack.com/p/efficient-extensible-expressive-typed). It has many extensibility advantages, which we'll take advantage of soon.

What does this have to do with automatic differentiation? AD is achieved by hard-coding how to differentiate primitive operations like addition and multiplication, and composing those primitive rules via the chain rule. The primitive operations define a language of tensor operations which we can interpret in a few ways - in particular, as straightforward tensor operations without differentiation via `Tensor`, as a forward mode differentiated program via `Forward`, or as a reverse mode differentiated program via `Reverse`. As for `RawTensor` we represent the primitive operations of the differentiable language as a trait, `Diffable`, and then implement this trait for each interpreter.

Let's start with the trait definition:

```rust
pub trait Diffable {
    type Elem: Num;

    fn log(&self) -> Self;
    fn exp(&self) -> Self;

    fn elementwise_add(&self, other: &Self) -> Self;
    fn elementwise_sub(&self, other: &Self) -> Self;
    fn elementwise_mul(&self, other: &Self) -> Self;
    fn elementwise_div(&self, other: &Self) -> Self;
    fn elementwise_pow(&self, exp: &Self) -> Self;
    fn elementwise_eq(&self, other: &Self) -> Self;

    fn sum(&self, axes: &[usize]) -> Self;
    fn max(&self, axes: &[usize]) -> Self;

    fn reshape(&self, shape: &[usize]) -> Self;
    fn permute(&self, dims: &[usize]) -> Self;
    fn expand(&self, shape: &[usize]) -> Self;
    fn pad(&self, padding: &[(usize, usize)]) -> Self;
    fn crop(&self, limits: &[(usize, usize)]) -> Self;

    fn new(shape: &[usize], data: &[Self::Elem]) -> Self;
    fn shape(&self) -> &[usize];
}
```

`Diffable`'s operations are similar to `RawTensor`'s, and we can categorize them in much the same way - unary operations, binary operations, reduce-like operations, and shape-changing operations. Missing is the optimized fused multiply-add in `RawTensor`, which illustrates the difference in intent between `RawTensor` and `Diffable`. While we could make `RawTensor` differentiable, I'll now try to convince you we don't want to.

Fused multiply-add is an optimized operation that we need on the lowest level to have some hope of efficiency. It is likely that to make Tensorken more efficient, we'll need to add more special-purpose operations to better exploit hardware primitives, reduce memory usage, and so on.

We don't (necessarily) want to figure out how to differentiate those special-purpose operations - we'd like a small set of primitive operations, define their derivatives, and then compose those into higher-level operations. We then get derivatives of those higher-level operations for free, because differentiation is so beautifully composable. Separating `Diffable` from `RawTensor` allows us to add efficient, special-purpose operations to `RawTensor` without figuring out their derivatives. Vice versa, we can add operations to `Diffable` without having to change `RawTensor` and its implementations.

Before `Diffable`, we translated user-facing operations like matrix multiplication to `RawTensor` operations, which were interpreted by a concrete `RawTensor` like `CpuRawTensor`. Now we add another interpreter, `Diffable`, between the user-facing operations and `RawTensor`, which not only calculates the primal results but also derivatives. `Diffable` interpreters execute both primal and derivative calculations as `RawTensor` operations. That means we can combine all implementations of `Diffable` with all implementations of `RawTensor`. So we can do forward AD on the GPU, reverse AD on the CPU, or any other combination.

Let's make our way down the interpreter layers to see how this works in practice. We'll start with matrix multiplication and end up at `CpuRawTensor`.

Each of the sections that follow is one layer of the interpreter lasagne:

- High-level tensor operations like `matmul` are translated to `Diffable` operations like `sum` and `elementwise_mul`.

- A `Diffable` interpreter like `Forward` and `Reverse` translates primitive operations like `sum` and `elementwise_mul` to `RawTensor` operations, adding calculation of derivatives.

- A `RawTensor` interpreter like `CpuRawTensor` executes the operations on a particular device.

![a white plate topped with lasagna covered in spinach](https://images.unsplash.com/photo-1633436374961-09b92742047b?crop=entropy&cs=tinysrgb&fit=max&fm=jpg&ixid=M3wzMDAzMzh8MHwxfHNlYXJjaHwzM3x8bGFzYWduZXxlbnwwfHx8fDE3MDIxNTcwOTJ8MA&ixlib=rb-4.0.3&q=80&w=1080 "a white plate topped with lasagna covered in spinach")

_A nicely layered design. Is it lunchtime yet? Photo by Parnis Azimi on Unsplash_

### User-facing layer: matrix multiplication in terms of `Diffable`

Here is a sketch of `matmul`, omitting everything that is not an operation on `Diffable`:

```rust
pub trait DiffableExt: Diffable
{
    fn matmul(&self, other: &Self) -> Self {
        // preconditions, shape manipulation omitted
        // special cases omitted

        let l = self.reshape(&l_shape);
        // shape manipulation omitted
        let r = other
            .reshape(&r_shape)
            .transpose(r_shape.ndims() - 1, r_shape.ndims() - 2);

        // after multiply: [..., m, o, n]
        l.mul(&r)

        // after sum: [..., m, o, 1]
        let sum = prod.sum(&[prod.shape().ndims() - 1]);

        // after reshape: [..., m, o]
        let s = sum.shape();
        sum.reshape(&s[..s.ndims() - 1])
    }
}

```

Tensorken has three implementations of `Diffable`: `Tensor`, `Forward`, and `Reverse`. `Tensor` doesn't do any differentiation at all - it translates `Diffable` to `RawTensor` operations. `Forward` and `Reverse` augment the operations with their respective mode of AD. We'll come back to these later - first, we need to find a Rust vehicle to put the user-facing operations that are not in `Diffable`. We could re-implement them on each implementation of `Diffable`, but that is redundant. Instead, I've defined `DiffableExt`, a sub-trait of `Diffable` with a blanket implementation:

```rust
pub trait DiffableExt: Diffable
{
    // all the fns we want, like matmul, go here.
    // They'll need to be defined in terms of Diffable,
    // because that's all that's available.

    fn matmul(&self, other: &Self) -> Self { ... }
}

impl<T: Diffable> DiffableExt for T {}

```

The advantage is we only have to implement `Diffable` on a concrete type, then anything defined on `DiffableExt` is available too (as long as `DiffableExt` is in scope.)

### **The first** `Diffable` **implementation: Tensor**

We now need a concrete type to present to users. `Tensor` is that type. Its definition is mysteriously simple:

```rust
pub struct Tensor<T>(T);
```

The idea is that the generic type argument `T` is a `Diffable`. Why not add the type constraint here? Because it's unnecessary - for all interesting implementations, `T` is `Diffable`. Constraining `T` here adds nothing new.

We can now make `Tensor<T>` implement `Diffable` for any `T` that's `Diffable`:

```rust
impl<T: Diffable> Diffable for Tensor<T> {
    type Elem = T::Elem;

    fn log(&self) -> Self {
        Tensor(self.0.log())
    }

    // etc
}
```

All operations delegate to `T`. Full implementation [here](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/tensor.rs#L160).

### From `Diffable` to `RawTensor`

That gets us nowhere - we can have a differentiable `Tensor<T>` if we have a differentiable `T`. To execute tensor operations we need to get to a `RawTensor`. We can do that by interpreting `Diffable` operations as `RawTensor` operations. In Rust, this means creating a [blanket implementation](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/tensor.rs#L22) of `Diffable` for any `RawTensor`:

```rust
impl<T: Num, TTensor: RawTensor<Elem = T>> Diffable for TTensor {
    type Elem = T;

    fn log(&self) -> Self {
        self.log()
    }

    // etc
}
```

Since `Diffable` is a subset of `RawTensor`, the implementation is again straightforward. A type like `Tensor<CpuRawTensor>` now works, and we can apply all operations in `Diffable` and `DiffableExt` to it.

It seems like we went around in a big circle. After Tensorken parts 1 and 2, we had a `Tensor<T: RawTensor>` with high-level operations like `matmul` and primitive operations on `RawTensor`. Now we have `Tensor<T: Diffable>` with high-level operations like `matmul` moved to `DiffableExt`, differentiable primitive operations on `Diffable`, and primitive executable operations still on `RawTensor`.

What we gained is the ability to have other `Diffable` implementations. We're going to use that ability now.

### Forward-mode AD with `Forward`

The `Forward` type wraps `T` with extra stuff so we can transform and trace the computation to calculate the derivative. In interpreter terms, `Forward` is an interpreter for the `Diffable` language that calculates the derivative alongside the primal result using forward-mode AD. It does that by applying all tensor operations on a dual tensor.

The `Forward` type:

```rust
pub enum Forward<T> {
    Lift(T),
    Forward(T, T),
}
```

Like for `Tensor<T>`, the `T` here is a `Diffable` tensor. The `Forward` case should make sense - it's the primal and the tangent tensors. We use the `Lift` case if we're not interested in computing the derivative of a tensor. Lifted tensors are treated as constants for the derivative computation. Another way of saying this is that their derivative is zero. We avoid many multiplications with zero by having a dedicated case instead of using the functionally equivalent `Forward(t, zero)`.

We can understand `jvp1`'s implementation now:

```rust
pub fn jvp1<T: Diffable + Clone, F>(f: F, at: &T, tangent: &T) -> (T, T)
where
    for<'a> F: Fn(&'a Forward<T>) -> Forward<T>,
{
    let forward = Forward::Forward(at.clone(), tangent.clone());
    let result = f(&forward);

    match result {
        Forward::Lift(p) => (p.clone(), p.zeros_like()),
        Forward::Forward(p, t) => (p, t),
    }
}
```

We wrap the `at` and `tangent` arguments in `Forward`, then call `f` with them and unwrap the `Forward` from the result.

`Forward` must implement `Diffable` for this to work. Finally, we come to the implementation of differentiation rules for the primitive operations:

```rust
impl<T: Clone + Diffable> Diffable for Forward<T> {
    type Elem = T::Elem;

    fn elementwise_mul(&self, rhs: &Self) -> Self {
        self.binary::<MulOp<T>>(rhs)
    }

    fn sum(&self, axes: &[usize]) -> Self {
        self.unary::<SumOp, _>(axes)
    }

    // etc
}
```

Full implementation [here](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/ad_forward.rs#L77).

Calculating the primal and derivatives are encapsulated in `Op` structs. The `unary` and `binary` functions deal with handling `Lift` or `Forward` enum cases in one place, and delegate to a given `Op` struct for the calculation:

```rust
impl<T: Diffable> Forward<T> {
    fn unary<Op: UnaryOp<T, Args = TArgs> + UnaryDiffOp<T>, TArgs: ?Sized>(
        &self,
        args: &TArgs,
    ) -> Self {
        let (primal, op) = Op::f(self.primal(), args);
        match self {
            Forward::Lift(_) => Forward::Lift(primal),
            Forward::Forward(_, tan) => Self::Forward(primal, op.dfda(tan)),
        }
    }
}
```

`binary` is [similar](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/ad_forward.rs#L64) but more involved because it has 4 combinations of `Lift` and `Forward`.

Here's `MulOp`:

```rust
pub(crate) struct MulOp<TTensor>(TTensor, TTensor);

impl<TTensor: Clone + Diffable> BinaryOp<TTensor> for MulOp<TTensor> {
    fn f(a: &TTensor, b: &TTensor) -> (TTensor, Self) {
        (a.elementwise_mul(b), MulOp(a.clone(), b.clone()))
    }
}

impl<TTensor: Diffable> BinaryDiffOp<TTensor> for MulOp<TTensor> {
    fn dfda(&self, d: &TTensor) -> TTensor {
        d.elementwise_mul(&self.1) // da * b
    }

    fn dfdb(&self, d: &TTensor) -> TTensor {
        d.elementwise_mul(&self.0) // db * a
    }
}
```

Differentiation rules often capture intermediate results or arguments of the primal computation. So `f` returns not only the result of the primal computation but also a struct to store whatever data is needed for the derivative computation. For `MulOp`, it captures the input tensors `a` and `b`.

`dfda` and `dfdb` define how to compute the derivative with respect to the first and second argument, given `d`, the derivative of downstream functions. The differentiation rule for elementwise tensor multiplication is essentially the same as for scalar multiplication.

Unary operations are similar but don't define `dfdb`:

```rust
pub(crate) struct SumOp(Vec<usize>);

impl<TTensor: Diffable> UnaryOp<TTensor> for SumOp {
    type Args = [usize];
    fn f(a: &TTensor, axes: &Self::Args) -> (TTensor, Self) {
        let r = a.sum(axes);
        (r, SumOp(axes.to_vec()))
    }
}

impl<TTensor: Diffable> UnaryDiffOp<TTensor> for SumOp {
    fn dfda(&self, d: &TTensor) -> TTensor {
        d.sum(&self.0)
    }
}
```

`SumOp` only needs the reduced axes from the primal computation to calculate `dfda` The derivative of the sum is the sum of the derivatives, so we can apply the same `sum` to primal and tangent.

You can find all the ops [here](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/ad_ops.rs) and [here](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/ad_forward.rs).

### `Forward<Forward<T>>` for higher order derivatives

Reiterating this signature:

```rust
pub fn jvp1<T: Diffable + Clone, F>(f: F, at: &T, tangent: &T) -> (T, T)
where
    for<'a> F: Fn(&'a Forward<T>) -> Forward<T>
```

Since the only requirement on `T` is that it's `Diffable` and `Forward` is `Diffable`, besides a `Tensor<T>` we can pass a `Forward<T>` to `jvp1` to calculate higher-order derivatives.

```
let p: Tensor<CpuRawTensor<f32>> = Tr::scalar(2.0);
let ddf = diff1(|x: &Forward<Tensor<_>>| 
            diff1(|x: &Forward<Forward<Tensor<_>>>| x.tanh(), x), 
            &p
        );
```

As we'll see soon, this same design will allow us to combine forward and reverse modes up to arbitrary depth, by building up types like `Forward<Reverse<Tensor<..>>>`.

This scheme works because the `Diffable` operations are implemented in terms of `Diffable`. That looks circular, but it's not: it's a stack of interpreters with a `Tensor` at the bottom, which translates `Diffable` operations to `RawTensor` operations:

`Forward<Forward<...>>: Diffable` ->... -> `Tensor: Diffable` -> `RawTensor`

Somewhat surprisingly, differentiating a differentiated program gets us the second derivative. One way to make sense of that is that if you compute second or third derivatives symbolically, that's exactly what you do: you apply the differentiation rules multiple times. If you want to do a "fun" exercise, you can work out that stacking `Forward` types amounts to operating on duals-of-duals up to the desired order. If you work out the differentiation rules by hand, you'll find that it yields the correct higher-order derivative.

### Reverse-mode AD with `Reverse`

The implementation of `Diffable` for `Reverse` follows the same pattern as `Forward` but is more involved. Because reverse mode accumulates the derivative in a separate backward pass, we can no longer compute everything on the fly when we compute the primal. Instead, `Reverse` builds a trace of operations in a forward pass while calculating the primal result, then accumulates derivatives in the backward pass.

The difference with forward mode is visible in the signature of `vjp`:

```rust
pub fn vjp1<'b, 't, T: Diffable + Clone + 't, F>(f: F, at: &T) -> (T, PullBack<'t, T>)
where
    for<'a> F: Fn(&'a Reverse<'a, 't, T>) -> Reverse<'a, 't, T>
```

Like `jvp`, it returns the primal result. Unlike `jvp`, it doesn't return the tangent, but instead a `PullBack` struct. The only available operation on that is `call`:

```rust
pub fn call(&self, cotangent: &T) -> Vec<T>
    where
        T: Diffable + Clone,
```

This takes a cotangent tensor - in other words, a tensor with the same shape as the result of `f`, and returns the tangents of all the arguments of `f`. Here's how `vjp` is used:

```rust
pub fn value_and_gradn<'t, T: Diffable + Clone + 't, F>(f: F, at: &[&T]) -> (T, Vec<T>)
where
    for<'a> F: Fn(&'a [Reverse<'a, 't, T>]) -> Reverse<'a, 't, T>,
{
    // one forward pass, tracing
    let (primal, pullback) = vjpn(f, at);
    // one backward pass, accumulating derivatives
    let tangents = pullback.call(&primal.ones_like());
    // but we get multiple tangents in one go
    (primal, tangents)
}
```

Other implementations for `grad` functions follow a similar pattern.

The details of how this is implemented (via a `Trace` type) are explained in [my post on AD](https://getcode.substack.com/p/automatic-differentiation-from-forward), so I won't repeat them here. It is not substantially different from the scalar case. Briefly, here is the `Reverse` type:

```rust
pub enum Reverse<'a, 't, T> {
    Lift(T),
    Reverse(&'a Trace<'t, T>, T, usize),
}
```

Like `Forward`, it has a `Lift` case for tensors we don't want to differentiate. The `Reverse` case contains the primal `T`, and some administrative data to record the trace and do the backward pass.

The implementation of `Diffable` looks similar to `Forward`:

```rust
impl<T: Clone + Diffable> Diffable for Reverse<'_, '_, T> {
    type Elem = T::Elem;

    fn elementwise_mul(&self, rhs: &Self) -> Self {
        self.binary::<MulOp<T>>(rhs)
    }

    fn sum(&self, axes: &[usize]) -> Self {
        self.unary::<SumOp, _>(axes)
    }

    // other omitted
}
```

Again we have `unary` and `binary` helper methods to deal with `Lift` and call the appropriate functions on the `Op` structs.

Interestingly, even though reverse mode calculates derivatives backward, from the output to the input, `MulOp` is identical for forward and reverse mode. This is true for all elementwise operations.

`sum` however is different from forward mode. In the backward pass, we get a `d` in the shape of the result of the `sum` (i.e. with fewer elements) and we need to produce a tensor in the shape of the input of `sum`. To do that, we need `expand`:

```rust
pub(crate) struct SumOp(Vec<usize>);

impl<TTensor: Diffable> UnaryOp<TTensor> for SumOp {
    type Args = [usize];
    fn f(a: &TTensor, axes: &Self::Args) -> (TTensor, Self) {
        let r = a.sum(axes);
        (r, SumOp(a.shape().to_vec()))
    }
}

impl<TTensor: Diffable> UnaryDiffOp<TTensor> for SumOp {
    fn dfda(&self, d: &TTensor) -> TTensor {
        d.expand(&self.0)
    }
}
```

Full implementation for reverse mode is in [ad\_reverse.rs](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/ad_reverse.rs) and the reverse operations are in [ad\_ops\_reverse.rs](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/ad_ops_reverse.rs).

After all that, we can run all the examples in the demo section. However, there is one remaining issue.

## Un-blowing up matmul, again

The problem is serious. Repeating the (pseudo-code) [implementation](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/diffable_ext.rs#L258) of `matmul` in `DiffableExt`:

```
pub trait DiffableExt: Diffable
{
    fn matmul(&self, other: &Self) -> Self {
        // preconditions, shape manipulation omitted
        // special cases omitted

        let l = self.reshape(&l_shape);
        // shape manipulation omitted
        let r = other
            .reshape(&r_shape)
            .transpose(r_shape.ndims() - 1, r_shape.ndims() - 2);

        // TROUBLE BEGINS HERE
        // after multiply: [..., m, o, n]
        l.mul(&r)

        // after sum: [..., m, o, 1]
        let sum = prod.sum(&[prod.shape().ndims() - 1]);

        // after reshape: [..., m, o]
        let s = sum.shape();
        sum.reshape(&s[..s.ndims() - 1])
    }
}
```

See that `mul` followed by `sum`? In the second part of Tensors from Scratch, I [explained that this blows up memory](https://getcode.substack.com/i/133922155/towards-matrix-multiplication), to the point where this approach is utterly unscalable. The fused multiply-add function in `RawTensor` came to the rescue - we rewrote the separate `sum` and `mul` calls into one `l.fused_multiply_add(&r, dims)`, which made it efficient. Now we've regressed to the previous bad situation. What gives?

First, `Diffable` doesn't have `fused_multiply_add`, so we can't write the optimized version directly. We could add `fused_multiply_add` to `Diffable` as a primitive operation, but then we have to define a differentiation rule for it in the various modes. One of the main reasons for `Diffable`'s existence is to avoid that.

Second, while manually fusing `mul` and `sum` worked for this particular case, users may inadvertently write a `mul` followed by a `sum`, and fall into this trap themselves. Worse, while we're calculating derivatives by composing operations in forward or reverse mode, Tensorken itself may introduce a `mul` followed by a `sum`. Manually fusing all cases is not going to work. We need a better solution.

If we were writing a compiler, it'd be straightforward to go through the abstract syntax tree of tensor operations and transform any `l.mul(r).sum(axes)` into `l.fused_multiply_add(r, axes)`. Can we do a similar optimization here?

Let's think about what's happening from the perspective of interpreters. We have defined a language for writing differentiable programs using the trait `Diffable`. Everything we do with tensors - `matmul`, `crop`, `max`, `sigmoid` as well as getting derivatives, is eventually a program in terms of the operations on `Diffable`. We have three interpreters for `Diffable` - one that translates the differentiable program to `RawTensor` operations, and two that augment the differentiable program with forward or reverse mode AD. No matter how many times we stack `Diffable` on top of `Diffable`, eventually the program gets run via a `RawTensor` interpreter.

We only have concrete `RawTensor` interpreters so far - that calculate the results on CPU or GPU, or that print a string representing the result. But we can also write an interpreter that spits out a new, optimized `RawTensor` interpreter, with all `mul` + `sum` fused into `fused_multiply_add`.

This technique - which I didn't invent at all, to be clear - is introduced more gradually and gracefully in [my post on typed tagless final interpreters](https://getcode.substack.com/p/efficient-extensible-expressive-typed). Here I'll give a whirlwind tour of the implementation.

We'll use a type called `Fuse<T>`. `T` is the target optimized `RawTensor`. Whenever `mul` followed by `sum` is detected in the unoptimized, original `RawTensor`, `Fuse` rewrites the two operations to a fused equivalent.

```rust
enum FuseCtx {
    Sum(Vec<usize>),
    NotSum,
}

pub struct Fuse<T>(Rc<dyn Fn(&FuseCtx) -> T>);
```

The function from `FuseCtx` to the fused `T: RawTensor` is a factory function we'll build up while interpreting the original `RawTensor` as `Fuse<T>`. In other words, `Fuse<T>` interprets `RawTensor` as a function that given a `FuseCtx` produces an optimized `RawTensor`. It works in two passes. A first pass builds up the factory function, then a second pass to run the function and get a new `RawTensor`.

Since `Fuse` only needs to fuse multiply and sum operations, it delays the application of `sum`, and instead passes `Sum(axes)` to the continuation via `FuseCtx`. The continuation calls the delayed `sum` if it can't fuse or `fused_multiply_add` if it can. Here's the implementation of `mul` where fusing happens:

```rust
impl<TRaw: RawTensor + Clone + 'static> RawTensor for Fuse<TRaw> {
    type Elem = TRaw::Elem;

    fn mul(&self, other: &Self) -> Self {
        let f_lhs = Rc::clone(&self.0);
        let f_rhs = Rc::clone(&other.0);
        let nextctx = FuseCtx::NotSum;
        Fuse::new(move |ctx| match ctx {
            FuseCtx::Sum(axes) => f_lhs(&nextctx).fused_multiply_add(&f_rhs(&nextctx), axes),
            FuseCtx::NotSum => f_lhs(&nextctx).mul(&f_rhs(&nextctx)),
        })
    }
}
```

The context passed in the closure represents what the next operation is, from the perspective of the current operation. If it's `sum`, the `Sum` enum case, we fuse. If it's anything else, represented by `NotSum`, we know the operation has already been applied and we can't fuse. Since `mul` is not a `sum`, we pass `NotSum` as the next context.

Here is the implementation of `sum`:

```rust
fn sum(&self, axes: &[usize]) -> Self {
    let f = Rc::clone(&self.0);
    let my_axes = axes.to_vec();
    Fuse::new(move |ctx| match ctx {
        FuseCtx::Sum(sum_axes) => f(&FuseCtx::Sum(combine_axes(&my_axes, sum_axes))),
        FuseCtx::NotSum => f(&FuseCtx::Sum(my_axes.clone())),
    })
}
```

We do not apply `sum` straight away to the resulting interpreter. Instead, we pass `Sum` through to the next operation, so it gets a chance to fuse it. Any operations that don't fuse, need to apply the delayed `sum` if they get the `Sum` enum. We might as well fuse consecutive `sum` calls into one by combining axes, hence the first match arm.

Fusing happens in two passes: the first pass builds the `FuseCtx -> RawTensor` function. The second pass creates the optimized `RawTensor` by calling the function:

```rust
impl<T> Fuse<T> {
    fn run(&self) -> T {
        (self.0)(&FuseCtx::NotSum)
    }
}
```

[Link to full implementation of fusing.](https://github.com/kurtschelfthout/tensorken/blob/v0.3/src/raw_tensor_fuse.rs)

Now I can finally reveal the full `Cpu32` and `Wgpu32` types:

```rust
pub type Cpu32 = Tensor<ShapeTracker<Fuse<CpuRawTensor<f32>>>>;
pub type Wgpu32<'d> = Tensor<ShapeTracker<Fuse<WgpuRawTensor<'d, f32>>>>;
```

The remaining unknown there is `ShapeTracker`. `ShapeTracker` is a `RawTensor` implementation that abstractly interprets the operations by only tracking tensor shapes. It delegates all operations to the `RawTensor` it wraps, except `shape`:

```rust
pub struct ShapeTracker<T>(ShapeStrider, T);

/// This implementation passes every operation through
/// to self.1, except for shape.
impl<TRaw: RawTensor> RawTensor for ShapeTracker<TRaw> {
    type Elem = TRaw::Elem;

    fn exp(&self) -> Self {
        Self(self.0.clone(), self.1.exp())
    }

    // etc

    fn shape(&self) -> &[usize] {
        self.0.shape()
    }
}
```

Because `Fuse` does not track shapes but does need to implement `RawTensor::shape`, it can only return its shape by running the delayed computation. We don't want that - some derivative operations require access to the shape of tensors, and it would be bad if we had to run the tensor computation at that point.

`ShapeTracker` solves this for us - it can answer `shape` queries without executing tensor operations. The order is important here. `ShapeTracker` needs to wrap `Fuse` which needs to wrap the concrete `CpuRawTensor` or `WgpuRawTensor`.

## I love it when a plan comes together

Thanks to the power of interpreters aka final tagless encoding, Tensorken gained a capable yet small and extensible AD implementation. So far, I'm really happy with how Tensorken turned out. I started seriously researching deep learning from an implementation perspective at the beginning of 2023 with only some prior exposure to automatic differentiation. I randomly ran into the [typed tagless final interpreters paper](https://www.cambridge.org/core/journals/journal-of-functional-programming/article/finally-tagless-partially-evaluated-tagless-staged-interpreters-for-simpler-typed-languages/7B2DC44A2127EBBA71ADE63809D9425F) while I was studying TinyGrad, and figured that TinyGrad's style would lend itself well to the tagless final style. I could not have hoped for a better outcome.

After that, I saw a post on Reddit praising JAX and immediately preferred the functional style over PyTorch's imperative AD interface. It was much more challenging to implement in Rust though! Those signatures look straightforward now, but it took a lot of struggling with closures, lifetimes and lifetimes and closures and then lifetimes some more before everything came together. All this to say - I got lucky trying an implementation style I hadn't ever used and struggled for a long time. When AD finally worked, it felt almost magical. Persistence is worth some IQ points.

Now that Tensorken has all the pieces of a full-fledged deep learning library, it's time to put it to the test. I intend to follow along with Andrej Karpathy's [Zero to Hero](https://karpathy.ai/zero-to-hero.html) neural networks course, translating it from PyTorch to Tensorken. At the end of that, we should have a home-grown, walking and talking nanoGPT. Without a doubt, there'll be many interesting problems in Tensorken itself to solve along the way.

Many thanks for reading!

## References

- [JAX: The Autodiff Cookbook](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html). I adapted some parts for this post.

- [JAX: Autodidax - JAX core from scratch](https://jax.readthedocs.io/en/latest/autodidax.html). Great insight into how JAX works under the hood, using a small implementation of JAX.

- [Video: Automatic Differentiation by Matthew Johnson]. Matthew Johnson is one of the authors of JAX. Here he talks about the principles underlying Autograd, another Python-based AD library.

- [TinyGrad](https://github.com/tinygrad/tinygrad). A small but powerful tensor library with reverse-mode AD in Python.

- [Swift's Differentiable Programming Manifesto](https://github.com/apple/swift/blob/main/docs/DifferentiableProgramming.md). Swift has a powerful differentiable programming component, integrated with the compiler.

- [Swift For Tensorflow (Google Drive)](https://drive.google.com/drive/folders/1y1FeJenXhoz_8xPw0ftVPngJUO_MdmRX). A great overview of Swift's approach to AD.

- [DiffSharp](https://diffsharp.github.io/). A tensor library with support for differentiable programming for .NET.