Thanks for the JAX pull request!
awesome.
1 Like
I've been trying to find out how to get this working since import _Differentiation doesn't seem to work but I can't seem to find any information on how to install it?
Try using any recent dev snapshot - it is not part of the release.
1 Like
Hi @Troy_Harvey, I've been experimenting with a complementary approach that can push things even further: graph compilation via tracing. The idea is to use Swift's existing autodiff as-is, but capture the computation into MLIR for XLA compilation.
Essentially, a DifferentiableTracer rides along with normal @differentiable code, observes Swift's AD generating pullbacks, emits MLIR operations, and hands the complete forward+backward graph to XLA for fusion into a single optimized kernel.
What's exciting is that this builds on top of the current autodiff—no compiler changes needed. The existing foundation is solid.
Repo here if curious: https://github.com/pedronahum/SwiftIR.
Worth noting: for the building simulation benchmark, at low iteration counts the Swift compiler and native execution is arguably impossible to beat. But as iterations scale up, the graph compilation approach starts to add real value—XLA's fusion and optimization have more to amortize.
New update posted here: Differentiable Swift Feb 2026 update
6 Likes