Performance & Compilation¶
The generated adjoint is ordinary Python — that's the product — but ordinary Python is also a performance ceiling. Tangent's answer has three parts: compile the adjoint, cache the compilation, and measure honestly.
Lower the adjoint: compile=¶
df = tangent.grad(f) # readable Python (default)
df = tangent.grad(f, compile='jax') # same adjoint, jax.jit-fused
df = tangent.grad(f, compile='torch') # torch.compile
df = tangent.grad(f, compile='tinygrad') # TinyJit
You inspect Python, you run compiled. The readable source stays attached
(df.__tangent_source__, verbose=1, tangent.explain),
and the cache stores the plain-Python function so every compile= choice
shares one differentiation.
Caveats: compile='jax' requires the primal to be written against
jax.numpy and traceable (fixed loops unroll; data-dependent Python control
flow does not trace). compile='torch' permits graph breaks, so tape
operations degrade gracefully. Numba is deliberately not offered — the
generated code calls Python-level tape helpers its nopython mode cannot
compile.
The compile bill, and the caches that pay it¶
Source transformation always pays a compile cost. Two caches amortize it:
- In-memory LRU — repeated
grad(f)calls in one process are ~µs. - Persistent disk cache — the generated source is stored under
~/.cache/tangent-ad/keyed by function source, options, and Tangent version, so new processes skip parsing, differentiation, and optimization entirely (measured ~13× faster cold start). Only source is stored, never pickles; on load, every global the compiled code references is verified to resolve, and any problem falls back to full recompilation. Disable withTANGENT_DISK_CACHE=0; relocate withTANGENT_CACHE_DIR.
Measured against the frameworks users would switch from¶
benchmarks/vs_frameworks.py, CPU (aarch64), best-of-50 eval; compile/first
call reported separately:
3-layer MLP loss (128×128, batch 64) — tensor-heavy:
| compile + first call | eval | |
|---|---|---|
| tangent (Python adjoint) | 32 ms | 0.181 ms |
tangent compile='jax' |
60 ms | 0.075 ms |
jax.grad + jit |
29 ms | 0.077 ms |
jax.grad (no jit) |
252 ms | 1.044 ms |
torch.autograd |
63 ms | 0.238 ms |
The lowered Tangent adjoint matches jax.grad+jit — the readable
source costs nothing at run time once compiled. Even the plain-Python adjoint
beats eager PyTorch and unjitted JAX here.
Scalar recurrence, 1000 data-dependent steps — control-flow-heavy:
| compile + first call | eval | |
|---|---|---|
| tangent (Python adjoint) | 61 ms | 3.5 ms |
torch.autograd (eager) |
34 ms | 22.2 ms |
jax.grad + jit (unrolled) |
6,802 ms | 0.017 ms |
This is the workload source transformation exists for: 6× faster than eager autograd per evaluation, at 1% of JAX's unroll-compile cost. (If you evaluate the same shapes thousands of times, JAX's unrolled kernel eventually wins — the table shows exactly where the crossover economics live.)
What the optimizer does¶
The generated code runs through tape-aware dead-code elimination, constant
folding, and assignment propagation by default (optimized=False to see the
naive adjoint), with optional SymPy-based algebraic simplification and CSE
(optimizations={'cse': True, 'algebraic': True}). DCE is tape-aware: a
push is only removed together with the pop that consumes it, so optimization
never corrupts higher-order derivatives.
Shrinking the tape: tape_liveness¶
Reverse mode saves each reassigned primal on the tape so the adjoint can read
it back. Many of those reads are shape-only: the value is only ever the
like of unbroadcast or the argument of init_grad, which consult its shape
and dtype but never its data. optimizations={'tape_liveness': True} stores a
lightweight shape carrier for exactly those entries instead of the whole array,
shrinking that tape slot from O(size) to O(ndim):
def loop(x):
s = x
for i in range(50):
a = s + 1.0 # `a` is read in reverse only for its shape
s = a * 0.5 + s * 0.5
return np.sum(s * s)
df = tangent.grad(loop, optimizations={'tape_liveness': True})
On a 20,000-element array over 50 iterations this cuts peak gradient memory by
~90%. The rewrite is proven safe by reaching-definition analysis - it converts
a variable only when every pop-reached use is shape-only, so any value the
adjoint reads arithmetically stays fully taped - and the gradient it produces
is identical to the default. It is opt-in (off by default) and targets the
NumPy execution path; it composes with √n checkpointing and compile=.
Fusing a straight-line segment: coarsening¶
For a straight-line stretch of elementwise arithmetic, optimizations={'coarsening': True}
differentiates the whole segment once symbolically (SymPy) and emits a single
vector-Jacobian-product expression per input instead of one adjoint statement
per primitive. Bound to the primal's backend, that one expression is handed to
JAX or PyTorch as a single fused kernel under compile=:
def f(x):
a = jnp.sin(x)
return jnp.sqrt(jnp.exp(a))
df = tangent.grad(f, optimizations={'coarsening': True}, compile='jax')
It is opt-in and applies only to coarsenable segments (elementwise straight-line code on NumPy/JAX/PyTorch); anything else - control flow, reductions, TensorFlow/Keras - falls back to the standard pipeline unchanged. See the coarsening deep-dive for the supported op set.