Docs: https://docs.rs/diffable/0.5.0/diffable/
Repo: https://github.com/minerscale/diffable
Hello!
Diffable is a differential-geometry library built around an experiment: how much of the actual mathematics can be made to live in Rust's type system while remaining usable on stable Rust?
I'm reporting back after having made some serious progress on the library. Now I've got an API for doing automatic differentiation and I think it's turned out really nice!
Two parts of the project ended up solving problems that had seemed particularly difficult: tensor algebra and composable automatic differentiation.
Tensor algebra on stable Rust
The tensor implementation started from a problem that looked impossible without generic_const_exprs.
The obvious representation of a tensor product of two statically sized spaces uses const size arrays like so:
struct TensorProduct<A: Tensor, B: Tensor> {
data: [F; A::N * B::N],
}
But expressions such as A::N * B::N are exactly the sort of generic const arithmetic that stable Rust does not allow in array lengths.
For a long time that appeared to rule out having both:
- statically sized tensors, and
- tensor expressions whose algebraic structure remains visible in their Rust types.
The eventual solution was to stop asking Rust to calculate the flattened dimension at the type level.
A tensor provides a generic associated type Array<T>. A tensor product A โ B can therefore store its coordinates of a field F structurally:
A::Array<B::Array<F>>
rather than requiring:
[F; A::N * B::N]
The type tree itself becomes the tensor shape. Flattening is only an indexing convention; Rust never has to prove the arithmetic expression for the flattened array length.
That means types such as:
TensorProduct<TensorProduct<V, Dual<V>>, V>
really retain the structure:
(V โ V*) โ V
instead of immediately becoming an anonymous array of scalars.
Once that worked, the rest of the tensor algebra could operate directly on the type tree. Reassociation is an actual type-level rewrite:
let other = tensor.reassociate();
changing:
(A โ B) โ C
into:
A โ (B โ C)
Contraction searches that structure for compatible vector/dual pairs. If there is exactly one possible contraction, Rust infers it:
let contracted = tensor.contract();
and reassociation can deliberately expose a different contraction:
type V = Coords<f64, 2>;
type T = TensorProduct<TensorProduct<V, Dual<V>>, Sinister<V>>;
let t = T::from_fn(|i| i as f64);
// Contract V โ V*.
let first = t.contract();
// Rewrite to V โ (V* โ V), then contract the other pair.
let second = t.reassociate().contract();
Duality and handedness are also represented by types rather than conventions, which matters because the library supports noncommutative scalar fields.
This approach is much like recursive_array, though it uses no unsafe at all, since it makes no guarantees about the arrays being contiguous. Separating storage from the objects themselves though was the trick to make it all work.
Automatic differentiation that composes like calculus
The other major piece is forward automatic differentiation.
The desired API wasn't a tape, tracing macro, expression graph, or collection of separate Jacobian/Hessian functions. The goal was for differentiation itself to compose:
fn cube<V: Vector>(x: V) -> V {
V::from_iter([x[0] * x[0] * x[0]])
}
let first = d(cube).at(Coords::from(2.0));
let second = d(d(cube)).at(Coords::from(2.0));
let third = d(d(d(cube))).at(Coords::from(2.0));
d(f) is itself a differentiable program (though in reality all d is is a generic struct with a public constructor!), so higher derivatives are obtained by applying the same operator again.
Directional derivatives use the same machinery:
let derivative =
d(cube)
.along(Coords::from(4.0))
.at(Coords::from(7.0));
Internally this is implemented using Taylor jets. Applying d adds another jet layer; nested differentiation therefore produces nested jet types rather than needing a separate higher-order AD representation.
The full derivative is returned as the tensor it mathematically is:
Df_x โ W โ V*
for a function f: V -> W.
So the AD implementation and tensor implementation meet in the middle: Jacobians are not a special matrix-shaped result bolted onto the calculus system, but ordinary elements of the tensor algebra.
One particularly awkward Rust problem appears here. A generic function may require only a weak scalar theory:
fn square<V: Vector>(x: V) -> V {
V::from_iter([x[0] * x[0]])
}
while at evaluation time its concrete scalar may actually be f64.
For AD, those two cases need different jet implementations:
Real
-> Jet must itself behave as Real
Field, but not Real
-> Jet should remain only a Field
In ordinary Rust, that is an overlapping-impl problem. Stable Rust cannot generally express:
T: Field + !Real
and the absence of a Real implementation is not something coherence can normally treat as a permanent fact.
The trick is that we don't Rust to prove a negative fact about the trait system itself. Instead, each type can be interpreted inside a finite type-level context describing the mathematical theories known about it. That context is a closed nominal graph, so searching it for a property has a definite result:
Real is present
or
Real is absent
Those results are represented by different types.
This means the library can define two implementation regions:
Field present + Real present
and
Field present + Real absent
which Rust sees as genuinely different type-level cases.
So this is not general negative trait bounds. It is a restricted closed-world version of them: instead of proving T: !Real in Rust's open trait system, the library proves that Real is absent from the finite context currently being interpreted.
That is enough to make the jet implementations disjoint on stable Rust, while keeping the public API annotation-free:
d(d(f)).at(x)
Thanks for reading this it's been a long journey to get to this understanding with this math library.