Provably Correct Tensor Puzzles

JAX is a numerical transformation system that utilizes a nearly one-to-one NumPy interface. While the primary use of JAX is machine learning, the core framework has remained impressively ML-agnostic. Instead, it focuses on a core set of transformations that can be reduced to simple numerical operations.

JAX is also a magic trick. The user thinks they are writing Python, while actually it is pulling out code using a minimal core of numerical operations. Proving properties of Python is nearly impossible, but proofs using this intermediate representation are more tractable.

If you are interested in the details of how JAX works, they have an amazing tutorial known as Autodidax. Here, we will just briefly outline the process.

Let's begin with a simple function that adds two vectors and multiplies the result by 2.

When we run this Python code, JAX is tracing the path and producing its own representation of what operations were called.

The intermediate representation is known as Jaxpr. It lives between your Python code and the final lower-level representations that execute the mathematical operations on hardware.

One day, we hope to have specs all the way down to the GPU. However, we are going to focus on this Jaxpr level and assume everything below roughly follows what it says. This is a horrible assumption for ML code, as numerical precision is of huge importance. Hopefully, we can cover this in future blogs.

Let's look at the Jaxpr output from our function. It should look familiar. It takes two arguments: float32 vectors of size 3. It adds them into a temporary variable, multiplies it by 2, and returns the result.

Of course, this code will get more complex with more realistic examples. There is a set of operations that specify the core numerical transformations needed. The core idea, though, is that it will remain a sequence of these operations, without loops or other control flow.

Our strategy to prove the correctness of this code is to transpile it over to Lean and then write proofs for the Lean version of the code. We do this with three Lean modules.

1. A mapping of the low-level Jaxpr
2. An implementation of a Jaxpr interpreter
3. A mapping of the high-level Jaxpr code.

If we trust our Jaxpr interpreter, we can then prove things about our JAX code in Lean and then prove that those properties also apply to the original Python code.

Let's walk through each of the steps.

First, we are going to map over the Jaxpr from the first section to Lean. Lean has a powerful macro system, so we can make the mapping look nearly identical to the original representation. We follow the specified syntax of the original, propagating forward some additional guarantees about bounds and dimensions.

Next, we need an implementation of the Jaxpr IR in Lean. As far as I know, there is not a formal specification of this language, but each part is well documented. In a more realistic setting, we would want to match the lower-level instructions exactly.

In addition to trusting that this interpreter matches Jaxpr, we are also going to relax the floating-point challenges at this stage. Our interpreter will just pretend things work like the reals.

Given the source of truth of the Jaxpr representation, we can now transpile the Jaxpr into something that is a bit more readable. This will look kind of like a line-aligned Lean implementation of the original code. Since this code will not be executed, we can apply transformations to make it a bit easier to use for our proofs (for instance, see the variance example below).

Since this implementation is generated, we shouldn't trust that it is correct. We therefore need to actually prove this. The following theorem shows that the code matches its IR.

If we can complete each of these three steps, we know that we have passed the implementation across the language barrier. Our implementation does this for most of the core Jaxpr IR with the assumption that we are working with the reals. From here on out, we can assume we are working with implementations in the proof language Lean.

We can now go back to the puzzles. We consider an example of the original outer puzzle, which asks the user to produce the outer product of two vectors. Here's the original code and the transpiled Lean.

In a similar way, we convert the loop version to Lean. This one is a bit long since JAX unrolls the loop to a sequence of instructions, so I am omitting the Lean code.

The goal is to certify that this produces the same output as the vectorized version. Once we are in Lean, proving this kind of thing is actually relatively straightforward. The code below automatically converts the broadcasting operations into expanded cases and enumerates all 2 × 3 cases to check that they produce the same result as the loop.

We combine this with the proof that the Jaxpr matches the above functions. We now have a full proof that the puzzle was correct!

Of course, just proving the correctness of puzzles is not that useful. Things become more interesting when we want to complete arbitrary ad hoc proofs about our code. To illustrate roughly what this might look like, let's walk through a few properties and look at:

• The JAX code
• The spec to prove
• An example of the proof (that an LLM might write for us)

Recently, I wrote a blog post on Lean-Verified Transformers that implemented several equivariance properties of transformers directly in Lean. We probably don't want to write the code in Lean, so we can write it in JAX and transpile it over.

We first define equivariance as:

The goal is to prove that a transformer without positional encoding is position-equivariant. We start with the transformer block in JAX.

A proof of part of this property would look like this.

And then we can propagate this through the whole stack.

Let's consider another simple, but common, example in ML. We want to guarantee that, after applying norm clipping in our code, we arrive at a gradient norm that is at most our fixed bound.

The goal here is to obtain the following certificate.

For a simple example, this is a mechanical calculation. But if you imagine an optimized implementation, we would need to push this through the whole neural network.

As a final example, we can consider the use of JAX for non-ML projects. The inspiration for this blog is the Noether project, which aimed to guarantee similar properties of mathematical code. The motivating example is conservation of properties through code.

The conservation law can be written with the following guarantee, which propagates through the system.

One thing that I have personally wanted for years is better guarantees for code that uses sampling, such as RL. Techniques like probabilistic programming provide these, but are not typically used in scaled systems. One solution might be to implement the JAX code as is and prove that it matches a clean spec for its probabilistic properties.

As a simple example, let us imagine doing Monte Carlo sampling that takes 16 samples of a die roll. We would like to establish the variance of our estimator.

The code itself will just be an imperative sampling implementation in JAX.

The Lean theorem will be like the math above.

In this case, the proof just propagates the variance through the operations.

As a final example, let's consider how these proofs might actually interact with all the work going on in formalizing mathematics. This is a toy example where we verify that the real-valued version of our JAX code produces the mathematically correct gradient.

For deeper examples of this, see the text on Verified Deep Learning with Lean 4.

We first define the function in JAX and then perform autodiff in the standard JAX manner to create a function for its derivative.

JAX will then lower the new function to Jaxpr that we can apply our method to. It goes through the conversion to Lean in the standard way.

There are many things we might prove here, but to keep things simple let's just get a certificate that the autodiff Jaxpr is equivalent to the mathematical derivative of the original function. We are using mathlib here with the deriv operator to produce that function directly.

This library provides a working implementation of the conversion between JAX and Lean with the ability to write non-trivial proofs. This is meant as a glance at what fully verifiable ML might allow for checking core properties like faithfulness, numerical stability, and statistical correctness. However, there is still lots to do here, including dealing with numerical precision, distributed components, making certificates more readable, and handling the devex of changing code and proofs. As in many areas, the future here will be quite unpredictable.