Rustograv

23 Aug, 2026

Turns out all you really need for automatic differentiation at compile time is to encode your tensor shenanigans, differential geometry as algebraic data types with the resulting type theory.

The derivatives aren't magically evaluated before the program runs. The point is that the structure of the mathematics can live in the types, and the compiler can then specialize the whole thing down into ordinary numerical code.

I've just finished the first working Rust version of autograv. The project started as a Python library using JAX for numerical relativity. It computes Christoffel symbols, curvature tensors, Ricci and Einstein tensors, stress-energy tensors, and the Kretschmann invariant from a metric function.

The Python version worked. But it was still very much an array-programming project: JAX arrays, jacfwd, einsum, and numerical tensor operations.

Then I came across u/minerscale's post about diffable 0.5.0 in r/rust.

diffable attacks the problem in two ways. The first was representing tensor products structurally instead of flattening them into one giant array.

The obvious representation of a tensor product looks something like:

[F; A::N * B::N]

The problem is that stable Rust (as of rustc 1.9.8) does not generally let you use arbitrary generic const arithmetic in array lengths.

diffable takes a different route. A tensor product can store its coordinates as a nested type:

A::Array<B::Array<F>>

The type tree becomes the tensor shape.

So instead of throwing away the mathematical structure and ending up with an anonymous flattened array, a type such as:

(V ⊗ V*) ⊗ V

can remain visible to the compiler as an actual tensor expression.

Reassociation and contraction can then operate on that tree directly.

The second idea was composable automatic differentiation through Taylor jets.

The desired interface isn't a separate Hessian function, a tracing macro, a tape, or an expression graph. Differentiation itself composes:

d(f)

d(d(f))

d(d(d(f)))

So each differentiation adds another jet layer. Higher derivatives are just nested jet types. That is a very different mental model from the way automatic differentiation is usually presented in application code.

I asked diffable's author if it could support index notation and tensor gymnastics. The answer was essentially: yes, but the difficult part is canonicalising all the different tensor trees that can represent the same index expression.

So I started sketching a Rust crate on top of diffable. The Rust side is not trying to recreate JAX. It uses diffable for the typed tensor algebra and Taylor-jet machinery, then adds the numerical-relativity layer on top.

The first awkward problem was getting a tensor-valued metric function through the jet machinery.

A metric is not just a scalar function. It maps coordinates to a rank-2 tensor. To differentiate it, the metric needs to evaluate once with ordinary floating-point coordinates and again with jet-valued coordinates, while preserving the tensor structure all the way through.

So I wrote a small tensor-valued Jacobian bridge that seeds the coordinate jets, evaluates the metric, and extracts the derivative tensor.

From there, the rest of the implementation became explicit tensor algebra:

• ⚡ A generic MetricField trait for metric functions over arbitrary scalar types

• 📐 Tensor views that expose rank-2, rank-3, and rank-4 diffable tensors as nested Rust arrays

• 🌌 Christoffel symbols, Torsion tensor, The Riemann cuvature tensor, The Ricci tensor and scalar curvature, the Einstein tensor, the stress-energy-momentum tensor, and the Kretschmann invariant computation.

The crate currently includes Minkowski, spherical-coordinate Euclidean, and Schwarzschild metrics.

The spherical-coordinate case was a nice reminder that coordinate complexity and physical curvature are different things. The metric has nonzero Christoffel symbols because spherical coordinates are curved as a coordinate system, but the Riemann tensor and Ricci scalar vanish because the underlying Euclidean space is flat.

The implementation uses no unsafe. The storage is deliberately boring. The interesting part is the type structure and the way the jet layers flow through the metric and curvature calculations.

There is also no final Einstein-notation macro yet. Writing a parser that lets you type something like R^a_bcd is not the hard part. The hard part is deciding how index expressions should canonicalise onto a tensor tree, how variance and handedness should be tracked, and how much of that should remain visible in the public API.

I would rather understand that problem properly than hide it behind a macro too early.

So this is the first Rust layer, not the final general-purpose tensor notation engine.

The nice part is that the mathematical core is now expressed in a form that Rust can reason about directly. The Python and Rust versions compute the same geometric quantities, but they have very different internal personalities.

Python/JAX gives you an extremely flexible array programming environment.

Rust plus diffable lets the tensor structure, differentiation structure, and algebraic constraints become part of the program's types.

This is also a good example of the kind of work I'm trying to do through Planckeon Labs.

The process was less about translating every line from Python into Rust and more about identifying the mathematical invariants, giving the compiler a precise target, reading the errors, testing each intermediate construction, and iterating until the implementation matched both the equations and the reference outputs.

The human role is choosing the problem, asking the right questions, and refusing to accept plausible-looking numerical output without verification.

The agent can do the exhaustive repository reading, compiler-driven iteration, test generation, documentation passes, and package checks.

tldr; the tensor tree becomes the intermediate representation, and the compiler becomes part of the mathematics.

Still early. Still experimental. But autograv now has a real Rust foundation.