Framework Gradient Benchmarks
Generated by benchmarks/vs_frameworks.py --emit. Do not edit by hand -
re-run the script so every number here comes from one source and cannot
drift.
- Environment: Python 3.12.3, Linux/aarch64, numpy 2.5.1, jax 0.11.1, torch 2.14.0+cpu
- Reported: best-of-50 steady-state eval time; compile+first call
(source transform / trace / jit) shown separately.
- Single machine, CPU. Absolute numbers vary by host; the ratios are
the point.
MLP loss (128x128, batch 64), d/dW1
| Framework |
eval (ms) |
compile+first (ms) |
notes |
| tangent (numpy) |
0.1813 |
2.1 |
|
| jax.grad (no jit) |
1.0863 |
316.3 |
|
| jax.grad + jit |
0.2575 |
40.9 |
|
| tangent compile='jax' |
0.2436 |
74.1 |
|
| torch.autograd |
0.3570 |
2.3 |
|
| torch.func.grad |
0.3847 |
483.2 |
|
| autograd (HIPS) |
0.2854 |
1.2 |
|
| finite differences |
3079.7 |
- |
projected 2x16384 params |
Conv layer im2col (16 imgs, 3->8 ch, 3x3 kernel), d/dK
| Framework |
eval (ms) |
compile+first (ms) |
notes |
| tangent (numpy) |
0.1909 |
1.3 |
|
| jax.grad + jit |
0.0887 |
31.2 |
|
| tangent compile='jax' |
0.1489 |
61.6 |
|
| torch.autograd |
0.2006 |
0.6 |
|
| torch.func.grad |
0.2502 |
0.4 |
|
| finite differences |
45.9860 |
- |
projected 2x216 params |
Scalar recurrence (1000 data-dependent steps), d/dx
| Framework |
eval (ms) |
compile+first (ms) |
notes |
| tangent (python) |
3.4692 |
5.7 |
|
| torch.autograd (eager) |
22.3970 |
23.5 |
|
| torch.func.grad (eager) |
42.0156 |
46.0 |
|
| autograd (HIPS) |
44.2050 |
161.3 |
|
| jax.grad + jit (unrolled) |
0.0216 |
6863.3 |
|
| finite differences |
0.0747 |
- |
2 forward evals |
How to read this
- Tensor-heavy (MLP, conv): generated NumPy is within ~2x of
jax.jit and beats eager frameworks; tangent compile='jax' matches jax.jit because it lowers the same adjoint through XLA.
- Loop / control-flow-heavy (recurrence): source transformation wins - it beats eager autograd/torch by several-fold, and
jax.jit only wins eval after paying a multi-second unroll compile.
- Finite differences are O(#params) and only approximate: cheapest on the single-input scalar case (two forward evals), but hopeless for the tensor cases (the projection shows why AD exists) and never exact.