Tangent - Source-to-Source Automatic Differentiation

A Python library for automatic differentiation that generates readable, inspectable gradient code — with NumPy, JAX, TensorFlow, PyTorch, Keras 3, and tinygrad support.

Originally developed by Google Research, now maintained and enhanced by @pedronahum.

Visualize how gradients flow through your computations

🌟 What Makes Tangent Unique?

Tangent performs source-to-source automatic differentiation: it transforms your Python code directly into gradient code that you can read, debug, and understand. Unlike black-box autodiff libraries:

• 📖 Readable: Generated gradient code is pure Python you can inspect
• 🔍 Debuggable: tangent.explain(f, x) shows the primal, the generated adjoint, a finite-difference check, the data-flow graph, and which inputs never affect the output — and tangent.source_map traces each adjoint line back to your code
• ✂️ Gradient surgery: read and edit the backward pass with with tangent.insert_grad_of(y) as dy: — scale, clip, log, or guard gradients mid-flow, stable across optimization passes (the demo no tracing AD can copy)
• ⌨️ Typed: ships py.typed (PEP 561), so df = tangent.grad(f) is a typed callable — mypy/pyright infer the gradient's signature, and for a single-argument function the gradient carries the input's type
• 🔧 Flexible: One API across NumPy, JAX, TensorFlow, PyTorch, Keras 3, and tinygrad
• 🐍 Pythonic: Control flow (break/continue/early return), closures, classes, list building, comprehensions, and second and third derivatives
• ⚡ Compilable: grad(f, compile='jax') runs the readable adjoint at jax.grad+jit speed; a persistent disk cache amortizes compilation across processes
• 🧩 Extensible: register custom gradients for your own functions or black-box ops (@tangent.custom_vjp, tangent.register_adjoint/register_tangent), plus stop_gradient and with tangent.checkpoint():
• 🔬 Differentiable simulators: tangent.odeint differentiates ODE solutions via the adjoint method (constant memory in step count)
• 🛡️ Shape-checked: tangent.check_shapes(f, x) catches rank/broadcast/matmul bugs at compile time, pointing at the offending line
• 📓 Works anywhere: @tangent.function captures source so gradients work in notebooks, the REPL, and exec-defined code; python -m tangent doctor diagnoses setup (and detects the dead 2017 tangent package)

import tangent

def f(x):
    return x ** 3 - 2 * x ** 2 + 3 * x - 1

df = tangent.grad(f, verbose=1)   # prints the generated gradient code
print(df(2.0))                    # f'(2) = 7.0

🚀 Quick Start

Installation

This fork is distributed on PyPI as tangent-ad (the import name stays import tangent, like Pillow's pillow/PIL). PyPI's tangent package is Google's unmaintained 2017 release — if you have it installed, run pip uninstall tangent first, since both provide the tangent module.

# Core (NumPy gradients)
pip install tangent-ad

# Optional backends are installed as extras:
pip install "tangent-ad[jax]"       # JAX support
pip install "tangent-ad[tf]"        # TensorFlow support
pip install "tangent-ad[torch]"     # PyTorch support
pip install "tangent-ad[keras]"     # Keras 3 (backend-agnostic keras.ops)
pip install "tangent-ad[tinygrad]"  # tinygrad (method-based tensor API)
pip install "tangent-ad[viz]"       # matplotlib + networkx visualization
pip install "tangent-ad[symbolic]"  # SymPy-based algebraic optimizations
pip install "tangent-ad[all]"       # everything above, plus the test tooling

# Or straight from the repository:
pip install "tangent-ad @ git+https://github.com/pedronahum/tangent.git"

Python 3.9–3.13 are supported and tested in CI. Note the platform caveats in Backend Support (e.g. TensorFlow has no CUDA build for aarch64).

Basic Usage

import tangent
import numpy as np

def f(x):
    return x ** 3 - 2 * x ** 2 + 3 * x - 1

df = tangent.grad(f)
print(df(2.0))   # 7.0

Multi-Backend Usage

The same tangent.grad API differentiates code written against any supported array library:

# NumPy
def f_np(x):
    return np.sum(x ** 2)

# JAX
def f_jax(x):
    return jnp.sum(jnp.tanh(x))

# TensorFlow
def f_tf(x):
    return tf.reduce_sum(tf.tanh(x))

# PyTorch
def f_torch(x):
    return torch.sum(torch.tanh(x))

# Keras 3 — works with whichever Keras backend is active
def f_keras(x):
    return kops.sum(kops.tanh(x))

# tinygrad — method-based tensor API
def f_tinygrad(x):
    return x.tanh().sum()

for f in (f_np, f_jax, f_tf, f_torch, f_keras, f_tinygrad):
    df = tangent.grad(f)

Each backend extension registers adjoints (reverse mode) and, for most ops, forward-mode tangents; see the per-extension lists in Backend Support.

🖥️ Backend Support

Backend notes worth knowing:

• Optional by design — extensions load when the backend is importable and warn (not fail) otherwise. Core NumPy autodiff always works.
• Keras backend choice — keras.ops gradients run on whichever backend Keras is configured with (KERAS_BACKEND=tensorflow|jax|torch). Verified with the TensorFlow and JAX backends; the torch backend works in torch-only environments, but importing TensorFlow and torch in one process can crash on some platforms (e.g. aarch64/GB10) — a platform issue, not a Tangent one.
• TensorFlow seeds — pass an explicit seed matching your tensor dtype (e.g. df(x, tf.constant(1.0, dtype=x.dtype))); TF does not implicitly mix float64 Python-float seeds with float32 tensors.
• Aliases — torch aliases (torch.negative/subtract/multiply/divide) are registered alongside the short names.
• tinygrad's method API — tinygrad exposes ops as methods (x.relu(), x.sum()). Tangent resolves calls against a function's globals, so a method on a computed tensor isn't statically resolvable; the tinygrad extension registers a method resolver that rewrites such calls to unbound form (Tensor.sum(x)). This is scoped to functions whose module actually imports tinygrad, so plain NumPy .sum() elsewhere is unaffected. The generated gradient code is itself a tinygrad graph, which tinygrad's compiler schedules and fuses.
• tinygrad conv/pool limitations — conv2d gradients require groups=1 (the input gradient supports any stride/dilation/padding; the weight gradient any stride/dilation); max_pool2d gradients require dilation=1, no ceil_mode, no return_indices, 3-D/4-D inputs, and symmetric padding; avg_pool2d gradients require the defaults (count_include_pad, no ceil_mode). Unsupported configurations raise NotImplementedError.
• The @ operator — x @ y differentiates for NumPy, JAX, TensorFlow, PyTorch and tinygrad (backend-dispatched via tangent.utils.register_matmul_grad).

A shared cross-backend test suite (tests/test_backend_coverage.py) runs one op catalog — arithmetic, exp/log, trig, activations, reductions, matmul, reshape, transpose — through every installed backend and checks the results against analytic gradients, keeping all six extensions at the same coverage bar.

🎨 Gallery of Gradients: See the Magic

The killer feature: readable gradient code.

Unlike black-box autodiff libraries, Tangent shows you exactly how gradients are computed. The gallery walks through 8 examples:

• 🔢 Polynomial derivatives (chain rule basics)
• 🔄 For loops that run in reverse during backprop
• 🌀 While loops with stack-based tape recording
• 🔀 Conditional branching (if/else and ternaries)
• 📊 NumPy array operations and broadcasting
• 📦 Nested function inlining
• 🔢 Matrix operations with colon slicing
• ⚡ Optimization comparison (before/after)

Each example shows: original function → generated gradient code → why it looks that way → verification.

→ Explore the Gallery | 📖 Documentation

🚀 Real-World Example: Building Energy Optimization

Optimize building heating to minimize energy costs — a real-world differentiable simulation.

Based on: PassiveLogic's Breaking the AI Speed Barrier and their Differentiable Swift Examples. This example demonstrates how Tangent achieves similar performance in Python.

import tangent

# Physical simulation: building temperature dynamics
def simulate_building(heating_schedule, outdoor_temp, electricity_price, params):
    T = params['T_initial']
    total_cost = 0.0
    for t in range(len(heating_schedule)):
        dT_dt = (outdoor_temp[t] - T) / (params['R'] * params['C']) + \
                heating_schedule[t] / params['C']
        T = T + params['dt'] * dT_dt
        total_cost += electricity_price[t] * heating_schedule[t] + \
                     params['lambda_comfort'] * (T - params['T_target']) ** 2
    return total_cost

# 🎯 Automatic differentiation — one line
grad_simulate = tangent.grad(simulate_building, optimized=True)
gradient = grad_simulate(heating_schedule, outdoor_temp, electricity_price, params)

→ Open Interactive Notebook

🐍 Python Feature Support

Tangent supports a broad subset of Python for numerical computing:

• Control flow: if/elif/else, for loops (with range/enumerate/zip), while loops, ternary expressions
• Operators: arithmetic, comparisons, boolean (and, or, not), augmented assignment (+=, -=, *=, /=, **=)
• Functions: lambdas, closures and factories, nested calls, default/keyword arguments
• Classes: user-defined classes with method inlining, instance attributes, method chaining, inheritance and super()
• Data: NumPy arrays; pytrees — tuples, lists and (nested) dicts of arrays can be passed as arguments and indexed/looped, with gradients returned in the same structure, and returned as outputs (the default seed is a matching pytree of ones — the gradient of the sum of all leaves — and a caller-supplied seed of the same structure is used as the cotangent); differentiable list building with xs.append(v) / v = xs.pop(); list comprehensions over both constant iterables (unrolled) and runtime iterables (lowered into differentiable loops, runtime if filters and nesting included); set/dict comprehensions over compile-time-constant iterables
• Statements: assert, pass, early return
• Higher-order: grad(grad(f)) second derivatives, third derivatives for NumPy functions, Hessian-vector products

📖 Complete reference with examples: Python Feature Support Guide — plus focused feature docs under docs/features/.

Supported: break/continue (lowered into guard flags with exact gradients). Not supported: @property/@classmethod/@staticmethod, generators, and in-place array mutation through augmented subscript assignment. See the feature guide for the full list.

⚡ Advanced Optimization Pipeline

tangent.grad(f, optimized=True) runs a multi-pass optimization pipeline that produces production-grade gradient code:

1. Constant folding — evaluate constant expressions at compile time
2. Dead code elimination — activity analysis + backward slicing remove unused computation (typically 30–50% of generated code); tape-aware, so push/pop pairs are only removed together
3. Assignment propagation — inline single-use variables
4. Strength reduction (opt-in) — x ** 2 → x * x, x / c → x * (1/c)
5. Common subexpression elimination (opt-in) — reuse repeated subexpressions
6. Algebraic simplification (opt-in) — SymPy-based identities (sin² + cos² → 1)
7. Fixed-point iteration — repeat until stable

Measured impact on the building-energy example: ~2.6× end-to-end, essentially all of it from tape-aware DCE. On this array-dominated workload the symbolic passes (strength reduction, CSE, algebraic) add nothing measurable — they help expression-heavy scalar code, which is why they are opt-in. See the Building Simulation Benchmark for the reproducible numbers.

df = tangent.grad(f, optimized=True)          # production (constant folding + DCE + propagation)
df = tangent.grad(f, optimized=False)         # all intermediate steps (education/debugging)
df = tangent.grad(f, optimized=True, verbose=1)  # prints what each pass did
df = tangent.grad(f, optimizations={'tape_liveness': True})  # store only the shape of primals the adjoint reads for shape (opt-in; big memory cut on array loops)

Tape-liveness (opt-in). Reverse mode saves each reassigned primal on the tape, but many are read back only for shape (as the target of unbroadcast/init_grad). optimizations={'tape_liveness': True} stores a lightweight shape carrier for exactly those, shrinking that tape slot from O(size) to O(ndim) — ~90% less peak gradient memory on a 20k-element array loop, with an identical gradient (proven safe by reaching-definition analysis).

Deep dives: Symbolic Optimizations · Strength Reduction · Performance Analysis · Straight-Line Coarsening · DCE implementation

Straight-line coarsening (opt-in)

A different, symbolic strategy is available for hot straight-line kernels. Instead of emitting one adjoint statement (and one tape push/pop) per primitive op, Tangent lifts the whole segment into a single SymPy expression, differentiates it once, and emits a compact vector-Jacobian product:

df = tangent.grad(f, optimizations={'coarsening': True})

For a small straight-line kernel the difference is dramatic — the per-op reverse pass emits dozens of adjoint statements and tape pushes, while coarsening emits one compact VJP:

def kernel(a, b, c):
    return (np.exp(np.sin(a * b)) + c) * a

# Coarsened adjoint (auto-generated):
#   ba = bret * (c + a*b*cos(a*b)*exp(sin(a*b)) + exp(sin(a*b)))
#   bb = bret * a**2 * cos(a*b) * exp(sin(a*b))
#   bc = a * bret

It currently coarsens the elementwise ops sin, cos, tan, exp, log, sqrt, arcsin, arccos, arctan (and + - * / **). Backend kernel handoff: the coarsened VJP is emitted in the backend the primal is written against — for a JAX or PyTorch primal it becomes one expression in that backend's ops, so pairing it with compile='jax' / torch.compile fuses the whole segment into a single kernel:

df = tangent.grad(jax_kernel, optimizations={'coarsening': True}, compile='jax')

It is a prototype and deliberately conservative: it only applies to reverse-mode gradients of pure straight-line elementwise segments. Anything else — control flow, reductions such as np.sum, TensorFlow/Keras primals, varargs, multi-output configurations, or preserve_result — transparently falls back to the standard pipeline, so enabling it never changes correctness. Requires the symbolic extra (see Installation). See tangent/optimizations/coarsening.py and docs/optimizations/COARSENING.md.

🎨 Visualization Tools

Interactive tools for understanding autodiff (install with the viz extra, see Installation):

import tangent
import matplotlib.pyplot as plt

def f(x):
    return x * x + 2.0 * x + 1.0

fig = tangent.visualize(f, mode='flow', inputs=(2.0,))
plt.show()

Run python examples/demo_visualization.py to generate the full set of demo plots.

🔬 Advanced Features

Multiple gradients

def f(x, y):
    return x * x * y + x * y * y

df = tangent.grad(f, wrt=(0, 1))
grad_x, grad_y = df(2.0, 3.0)   # ∂f/∂x = 21.0, ∂f/∂y = 16.0

Preserve results

df = tangent.grad(f, preserve_result=True)
gradient, result = df(x)   # both the gradient and the primal value

Forward and reverse mode

df_rev = tangent.grad(f)                          # reverse mode (default)
df_fwd = tangent.autodiff(f, mode='forward')      # forward mode

Second and third derivatives

def f(x):
    return x ** 3

ddf = tangent.grad(tangent.grad(f))
print(ddf(2.0))   # 12.0  (d²/dx² x³ = 6x)

dddf = tangent.grad(ddf)
print(dddf(2.0))  # 6.0   (d³/dx³ x³ = 6)

Second derivatives work in both optimized and unoptimized modes, including through loops, across backends, and through hand-written use of the low-level tape API (tangent.push/pop/Stack) — dead-code elimination is tape-aware and only removes push/pop pairs together. Third derivatives (differentiating the second-order adjoint code again) work for ordinary NumPy functions in both modes. Fourth order and above is not yet reliable.

Custom gradients

Supply a gradient Tangent can't (or shouldn't) derive — for a numerically better rule, or to wrap a black box it can't transform:

@tangent.custom_vjp                 # your own function / a black box you call
def gelu(x): ...
@gelu.defvjp
def gelu_vjp(g, ans, x): ...        # plain-Python reverse rule

@tangent.register_adjoint(np.hypot)  # attach a rule to a library op you don't own
def _(z, x, y):
    d[x] = d[z] * x / z
    d[y] = d[z] * y / z

See the Custom Gradients guide.

Typed gradients

Tangent ships a py.typed marker, so df = tangent.grad(f) is a typed callable to mypy/pyright — no Any. For a single-argument function the gradient carries the input's type (grad(f: Callable[[T], Any]) -> Callable[..., T]).

Automatic caching

Compiled gradient functions are cached per (source, transform options), so repeated tangent.grad(f) calls are orders of magnitude faster than the first compilation. Inspect with tangent.get_cache_stats().

⚠️ Known Limitations

Documented honestly — see Python Feature Support for details:

• Fourth-order and higher derivatives: second and third derivatives are supported, but the fourth and beyond are not yet reliable — the optimized path can return incorrect values and the unoptimized path reaches the tape machinery above.
• Containers (pytrees): arguments and return values may be tuples, lists and (nested) dicts of arrays, including second derivatives through container arguments. A structurally mismatched container seed is rejected with a ValueError (tangent.match_seed reconciles the seed with the return value's structure at runtime). Third and higher derivatives through container outputs are untested territory.
• TF seed dtype — see Backend Support.
• SciPy support (scipy.special, scipy.linalg.solve/inv) loads when SciPy is installed; see Backend Support.
• Gradient checkpointing (tangent.grad(f, checkpoint=True)) uses segment (√n) recomputation: eligible loops run untaped with state snapshots every ceil(√n) iterations, and the backward pass replays one segment at a time — peak tape memory drops from O(n) to O(√n) with identical gradients (measured: 49.2 MB → 1.6 MB, a 96.8% reduction, on a 2000-iteration loop carrying a 1000-float state; see benchmarks/checkpointing_memory.py). Eligible for loops (constant bounds ≥ 100) use segment checkpointing; while loops and data-dependent iteration use online (Stumm–Walther) checkpointing under a fixed snapshot budget; with tangent.checkpoint(): annotates specific loops (any runtime bound). First-order reverse mode only — higher-order derivatives of a checkpointed gradient are not supported. See the Checkpointing User Guide.

📖 Examples

Linear regression

import tangent
import numpy as np

X = np.random.randn(100, 1)
y = 3 * X + 2 + np.random.randn(100, 1) * 0.5

def mse_loss(w, b):
    return np.mean((w * X + b - y) ** 2)

dmse_dw = tangent.grad(mse_loss, wrt=(0,))
dmse_db = tangent.grad(mse_loss, wrt=(1,))

w, b, lr = 0.0, 0.0, 0.1
for epoch in range(50):
    w -= lr * dmse_dw(w, b)
    b -= lr * dmse_db(w, b)

Neural network with JAX

import tangent
import jax.numpy as jnp
import jax

def neural_network(W1, b1, W2, b2, X, y):
    hidden = jax.nn.relu(jnp.dot(X, W1) + b1)
    output = jax.nn.sigmoid(jnp.dot(hidden, W2) + b2)
    return -jnp.mean(y * jnp.log(output) + (1 - y) * jnp.log(1 - output))

dnn_dW1 = tangent.grad(neural_network, wrt=(0,))
# ... one grad per parameter; use in your training loop

Runnable scripts

See examples/ for self-contained, runnable demos. The best starting point is examples/recent_features.py, which showcases coarsening, pytree arguments, multi-backend gradients, differentiable concat/stack, and second derivatives in one script:

python examples/recent_features.py

Notebooks

📓 Full notebook gallery → — every notebook, runnable in Colab, one obvious entry point.

• Tangent Tutorial — general introduction
• Gradient Surgery — read and edit the generated backward pass (insert_grad_of), the demo no tracing AD can copy
• LIBOR Market Model greeks — case study: differentiate a messy NumPy Monte-Carlo rate simulator for exact deltas/vegas in one reverse pass
• Projectile with drag — case study: sensitivities of a nonlinear ODE integrator, and optimizing the launch by gradient descent
• SIR epidemic calibration — case study: differentiate through a data-dependent branch to fit an epidemic model
• Gallery of Gradients — readable-code showcase
• Building Energy Optimization — real-world application

🧪 Testing

The suite covers core autodiff, every Python feature, all six backends, and second and third derivatives. The cross-backend catalog is checked two ways: against hand-derived analytic gradients and against an independent finite-difference oracle, so newly added adjoints are verified automatically:

pytest tests/                        # core suite (no backends required)
pytest tests/test_backend_coverage.py  # cross-backend op catalog (+ FD oracle)
pytest tests/test_coarsening.py      # straight-line coarsening
pytest tests/test_matmul_operator.py # @ operator (NumPy)
pytest tests/test_numpy_method_calls.py  # NumPy array-method rewrite
pytest tests/test_torch.py           # PyTorch-specific tests
pytest tests/test_keras.py           # Keras tests (any backend)
pytest tests/test_tinygrad.py        # tinygrad-specific tests

Current status: 79,000+ parameterized test cases pass, with zero failures and zero expected-failure (xfail) markers (last local measurement: 79,773 passed, 55 skipped, with NumPy, JAX, PyTorch-CPU, tinygrad, Keras 3, SymPy, SciPy and TensorFlow all installed). CI runs the suite on Python 3.9–3.13, exercises the cross-backend catalog against every installed backend, and enforces ruff check + ruff format (run make check locally, or make fmt to auto-format).

📊 Repository Structure

tangent/
├── tangent/                      # Core library
│   ├── grad_util.py              # Main autodiff engine (grad/autodiff/vjp/jvp)
│   ├── reverse_ad.py             # Reverse-mode transformation
│   ├── forward_ad.py             # Forward-mode transformation
│   ├── grads.py                  # Core NumPy adjoints
│   ├── tangents.py               # Core forward-mode tangents
│   ├── numpy_extended.py         # Extended NumPy adjoints
│   ├── jax_extensions.py         # JAX adjoints + tangents
│   ├── tf_extensions.py          # TensorFlow adjoints + tangents
│   ├── tf_extended.py            # Extended TensorFlow adjoints
│   ├── torch_extensions.py       # PyTorch adjoints + tangents
│   ├── keras_extensions.py       # Keras 3 (backend-agnostic) adjoints + tangents
│   ├── tinygrad_extensions.py    # tinygrad adjoints + tangents
│   ├── class_desugar.py          # Class method inlining
│   ├── visualization.py          # Visualization tools
│   ├── function_cache.py         # Gradient-function caching
│   ├── optimizations/            # DCE, CSE, strength reduction, algebraic
│   └── checkpointing_simple.py   # Manual forward-pass checkpointing helpers
├── tests/                        # 90+ test modules (79k+ cases)
│   ├── test_backend_coverage.py  # Cross-backend op catalog
│   ├── test_torch.py             # PyTorch tests
│   ├── test_keras.py             # Keras tests
│   ├── test_jax.py               # JAX tests
│   ├── test_tensorflow.py        # TensorFlow tests
│   └── ...                       # Feature and engine tests
├── examples/                     # Scripts and notebooks
├── notebooks/                    # Interactive tutorials
├── benchmarks/                   # Performance benchmarks
└── docs/                         # Feature/optimization/benchmark docs

📚 Documentation

• → Documentation site · Full index
• Tangent 2 writeup — a short, reproducible paper (runnable notebook + paper/reproduce.py)
• Notebook Gallery — every runnable Colab notebook, one entry point
• Python Feature Support — the definitive feature reference
• Custom Gradients — custom_vjp, register_adjoint/register_tangent, stop_gradient
• Performance & Compilation — compile=, tape-liveness, coarsening, and the honest framework benchmarks
• Checkpointing User Guide — gradient checkpointing: what works today and what doesn't

🤝 Contributing

Contributions are welcome! See CONTRIBUTING.md. Good starting points:

• Additional gradient definitions (extend any backend module + the cross-backend catalog)
• The open items in Known Limitations
• Documentation improvements

1. Fork the repository
2. Create a feature branch
3. Add tests for your changes (pytest tests/ must stay green)
4. Submit a pull request

📝 License

Apache License 2.0 — see LICENSE.

Original work Copyright 2017 Google Inc. Modified work Copyright 2024–2026 Pedro Nahum

🙏 Acknowledgments

• Original Tangent library by Google Research
• The JAX, TensorFlow, and PyTorch teams for the numerical computing ecosystems
• The Python scientific computing community

📬 Contact

• Repository: github.com/pedronahum/tangent
• Issues: github.com/pedronahum/tangent/issues
• Author: @pedronahum

Built with ❤️ for the machine learning and scientific computing communities