How autograd works

from pathlib import Path

from .graph import save_graph_html
from .tensor import Tensor


def _source_files() -> dict[str, str]:
    directory = Path(__file__).parent
    names = ("example.py", "tensor.py", "graph.py")
    return {name: (directory / name).read_text(encoding="utf-8") for name in names}


def main() -> None:
    x = Tensor(2.0, requires_grad=True, label="x")
    w = Tensor(-3.0, requires_grad=True, label="w")
    b = Tensor(10.0, requires_grad=True, label="b")

    # Python operators construct this graph as they execute:
    # x --\
    #      (*) --\
    # w --/       (+) --> ReLU --> loss
    # b ----------/
    prediction = (x * w + b).relu()
    loss = prediction**2
    loss.label = "loss"

    loss.backward()

    print(f"loss = {loss.data}")
    print(f"d(loss)/dx = {x.grad}")
    print(f"d(loss)/dw = {w.grad}")
    print(f"d(loss)/db = {b.grad}")

    output = save_graph_html(
        loss,
        Path.cwd() / "autograd_graph.html",
        source_files=_source_files(),
    )
    print(f"graph written to {output}")


if __name__ == "__main__":
    main()

from __future__ import annotations

import math
from collections.abc import Callable


Scalar = int | float


class Tensor:
    """A scalar value with a dynamically constructed autograd graph."""

    def __init__(
        self,
        data: Scalar,
        *,
        requires_grad: bool = False,
        label: str = "",
        _parents: tuple[Tensor, ...] = (),
        _op: str = "",
    ) -> None:
        self.data = float(data)
        self.requires_grad = requires_grad
        self.grad = 0.0
        self.label = label

        self._parents = _parents
        self._op = _op
        self._backward: Callable[[], None] = lambda: None

    def __repr__(self) -> str:
        return (
            f"Tensor(data={self.data:g}, grad={self.grad:g}, "
            f"requires_grad={self.requires_grad})"
        )

    @staticmethod
    def _coerce(value: Tensor | Scalar) -> Tensor:
        if isinstance(value, Tensor):
            return value
        return Tensor(value)

    def __add__(self, other: Tensor | Scalar) -> Tensor:
        other = self._coerce(other)
        out = Tensor(
            self.data + other.data,
            requires_grad=self.requires_grad or other.requires_grad,
            _parents=(self, other),
            _op="+",
        )

        def _backward() -> None:
            if self.requires_grad:
                self.grad += out.grad
            if other.requires_grad:
                other.grad += out.grad

        out._backward = _backward
        return out

    def __radd__(self, other: Scalar) -> Tensor:
        return self + other

    def __mul__(self, other: Tensor | Scalar) -> Tensor:
        other = self._coerce(other)
        out = Tensor(
            self.data * other.data,
            requires_grad=self.requires_grad or other.requires_grad,
            _parents=(self, other),
            _op="*",
        )

        def _backward() -> None:
            if self.requires_grad:
                self.grad += other.data * out.grad
            if other.requires_grad:
                other.grad += self.data * out.grad

        out._backward = _backward
        return out

    def __rmul__(self, other: Scalar) -> Tensor:
        return self * other

    def __neg__(self) -> Tensor:
        return self * -1

    def __sub__(self, other: Tensor | Scalar) -> Tensor:
        return self + -self._coerce(other)

    def __rsub__(self, other: Scalar) -> Tensor:
        return self._coerce(other) - self

    def __truediv__(self, other: Tensor | Scalar) -> Tensor:
        other = self._coerce(other)
        out = Tensor(
            self.data / other.data,
            requires_grad=self.requires_grad or other.requires_grad,
            _parents=(self, other),
            _op="/",
        )

        def _backward() -> None:
            if self.requires_grad:
                self.grad += out.grad / other.data
            if other.requires_grad:
                other.grad -= out.grad * self.data / (other.data**2)

        out._backward = _backward
        return out

    def __rtruediv__(self, other: Scalar) -> Tensor:
        return self._coerce(other) / self

    def __pow__(self, exponent: Scalar) -> Tensor:
        out = Tensor(
            self.data**exponent,
            requires_grad=self.requires_grad,
            _parents=(self,),
            _op=f"**{exponent:g}",
        )

        def _backward() -> None:
            if self.requires_grad:
                self.grad += exponent * self.data ** (exponent - 1) * out.grad

        out._backward = _backward
        return out

    def exp(self) -> Tensor:
        out = Tensor(
            math.exp(self.data),
            requires_grad=self.requires_grad,
            _parents=(self,),
            _op="exp",
        )

        def _backward() -> None:
            if self.requires_grad:
                self.grad += out.data * out.grad

        out._backward = _backward
        return out

    def tanh(self) -> Tensor:
        value = math.tanh(self.data)
        out = Tensor(
            value,
            requires_grad=self.requires_grad,
            _parents=(self,),
            _op="tanh",
        )

        def _backward() -> None:
            if self.requires_grad:
                self.grad += (1 - value**2) * out.grad

        out._backward = _backward
        return out

    def relu(self) -> Tensor:
        out = Tensor(
            max(0.0, self.data),
            requires_grad=self.requires_grad,
            _parents=(self,),
            _op="ReLU",
        )

        def _backward() -> None:
            if self.requires_grad:
                self.grad += (self.data > 0) * out.grad

        out._backward = _backward
        return out

    def _topological_order(self) -> list[Tensor]:
        order: list[Tensor] = []
        visited: set[Tensor] = set()

        def visit(node: Tensor) -> None:
            if node in visited:
                return
            visited.add(node)
            for parent in node._parents:
                visit(parent)
            order.append(node)

        visit(self)
        return order

    def backward(self, gradient: Scalar = 1.0) -> None:
        """Apply reverse-mode autodiff to this scalar output."""
        if not self.requires_grad:
            raise RuntimeError("cannot call backward() when requires_grad is False")

        order = self._topological_order()

        # Intermediates are working storage for this pass. Leaf gradients keep
        # accumulating, matching the behavior users see on PyTorch parameters.
        for node in order:
            if node._parents:
                node.grad = 0.0

        if self._parents:
            self.grad = float(gradient)
        else:
            self.grad += float(gradient)

        for node in reversed(order):
            node._backward()

    def zero_grad(self) -> None:
        """Clear gradients on every Tensor reachable from this Tensor."""
        for node in self._topological_order():
            node.grad = 0.0

from __future__ import annotations

import html
from collections.abc import Mapping
from pathlib import Path
from string import Template

from .tensor import Tensor


_BOX_WIDTH = 190
_BOX_HEIGHT = 82
_EXAMPLE_SNIPPET = """x = Tensor(2.0, requires_grad=True, label="x")
w = Tensor(-3.0, requires_grad=True, label="w")
b = Tensor(10.0, requires_grad=True, label="b")

prediction = (x * w + b).relu()
loss = prediction**2
loss.label = "loss"

loss.backward()
"""

_PAGE_TEMPLATE = Template("""<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>How autograd works</title>
  <style>
    :root { color-scheme: light; --ink: #172033; --muted: #62708a; --line: #dbe1ec; --blue: #2866c7; }
    * { box-sizing: border-box; }
    body { margin: 0; background: #f5f7fb; color: var(--ink); font: 14px system-ui, sans-serif; }
    header { max-width: 1440px; margin: auto; padding: 28px 28px 0; }
    h1 { margin: 0 0 7px; font-size: 26px; }
    h2 { margin: 0 0 14px; font-size: 18px; }
    p { margin: 0; color: var(--muted); }
    .tabs { display: flex; gap: 6px; max-width: 1440px; margin: 22px auto 0; padding: 0 28px; border-bottom: 1px solid var(--line); }
    .tab { border: 0; border-bottom: 3px solid transparent; padding: 11px 16px; background: transparent; color: var(--muted); font-weight: 650; cursor: pointer; }
    .tab.active { border-color: var(--blue); color: var(--blue); }
    .panel { max-width: 1440px; margin: auto; padding: 22px 28px 36px; }
    .card { overflow: hidden; margin-bottom: 22px; border: 1px solid var(--line); border-radius: 12px; background: white; box-shadow: 0 2px 8px #23304a0d; }
    .card-title { padding: 17px 20px 0; }
    .example-code { margin: 16px 20px 4px; padding: 16px 18px; overflow: auto; border: 1px solid #d9e2f1; border-radius: 9px; background: #f7f9fc; color: #24324a; font: 13px/1.55 ui-monospace, SFMono-Regular, Consolas, monospace; }
    .canvas { overflow: auto; }
    svg { display: block; min-width: 100%; }
    .edge { fill: none; stroke: #7b879c; stroke-width: 2; }
    .node rect { fill: #e7f0ff; stroke: #3973c6; stroke-width: 2; }
    .node.constant rect { fill: #f3f4f6; stroke: #9aa3b2; }
    .node text { fill: #28344a; font: 13px ui-monospace, monospace; }
    .node .name { font-weight: 700; }
    .node .op { fill: #8b3fb0; font-weight: 700; }
    .table-wrap { overflow: auto; padding-top: 8px; }
    table { width: 100%; border-collapse: collapse; }
    th, td { padding: 12px 16px; border-top: 1px solid #edf0f5; text-align: left; white-space: nowrap; }
    th { color: var(--muted); font-size: 12px; letter-spacing: .04em; text-transform: uppercase; }
    td code { color: #733a9c; }
    .step { width: 52px; color: var(--muted); }
    .code-browser { display: grid; grid-template-columns: 220px minmax(0, 1fr); min-height: 660px; overflow: hidden; border: 1px solid #2d3648; border-radius: 12px; background: #111722; box-shadow: 0 8px 28px #17203324; }
    .file-list { padding: 12px 8px; border-right: 1px solid #30394b; background: #18202d; }
    .file-button { display: block; width: 100%; border: 0; border-radius: 6px; padding: 9px 11px; background: transparent; color: #abb7ca; font: 13px ui-monospace, monospace; text-align: left; cursor: pointer; }
    .file-button:hover { background: #222d3d; color: white; }
    .file-button.active { background: #283b58; color: #8fc5ff; }
    .source-view { min-width: 0; overflow: auto; }
    .source-view pre { min-height: 100%; margin: 0; padding: 20px 24px; color: #d8dee9; font: 13px/1.55 ui-monospace, SFMono-Regular, Consolas, monospace; tab-size: 4; }
    .empty { padding: 28px; color: #abb7ca; }
    @media (max-width: 720px) {
      .code-browser { grid-template-columns: 1fr; }
      .file-list { display: flex; overflow: auto; border-right: 0; border-bottom: 1px solid #30394b; }
      .file-button { width: auto; white-space: nowrap; }
    }
  </style>
</head>
<body>
  <header>
    <h1>How autograd works</h1>
    <p>The forward pass constructs a graph; the backward pass applies the chain rule in reverse.</p>
  </header>
  <nav class="tabs" aria-label="Tutorial sections">
    <button class="tab active" type="button" data-tab-target="graph-panel" aria-selected="true">Graph &amp; walkthrough</button>
    <button class="tab" type="button" data-tab-target="code-panel" aria-selected="false">Code</button>
  </nav>
  <main>
    <section class="panel" id="graph-panel">
      <div class="card">
        <div class="card-title">
          <h2>Computation graph</h2>
          <p>Arrows run from each input to the value it produces. Blue nodes track gradients.</p>
        </div>
        <pre class="example-code"><code>$example_code</code></pre>
        <div class="canvas">
          <svg xmlns="http://www.w3.org/2000/svg" width="$width" height="$height" viewBox="0 0 $width $height">
            <defs>
              <marker id="arrow" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
                <polygon points="0 0, 10 3.5, 0 7" fill="#7b879c" />
              </marker>
            </defs>
            $edges
            $boxes
          </svg>
        </div>
      </div>
      <div class="card">
        <div class="card-title">
          <h2>Example walkthrough</h2>
          <p>Rows are in forward execution order. Gradients show the completed backward pass.</p>
        </div>
        <div class="table-wrap">
          <table>
            <thead><tr><th>Step</th><th>Tensor</th><th>Created by</th><th>Forward value</th><th>Backward gradient</th></tr></thead>
            <tbody>$walkthrough</tbody>
          </table>
        </div>
      </div>
    </section>
    <section class="panel" id="code-panel" hidden>
      <div class="code-browser">$source_browser</div>
    </section>
  </main>
  <script>
    const tabButtons = document.querySelectorAll("[data-tab-target]");
    const tabPanels = document.querySelectorAll("main > .panel");
    for (const button of tabButtons) {
      button.addEventListener("click", () => {
        for (const candidate of tabButtons) {
          const selected = candidate === button;
          candidate.classList.toggle("active", selected);
          candidate.setAttribute("aria-selected", String(selected));
        }
        for (const panel of tabPanels) {
          panel.hidden = panel.id !== button.dataset.tabTarget;
        }
      });
    }

    const fileButtons = document.querySelectorAll("[data-source-target]");
    const sourcePanels = document.querySelectorAll("[data-source-panel]");
    for (const button of fileButtons) {
      button.addEventListener("click", () => {
        for (const candidate of fileButtons) {
          candidate.classList.toggle("active", candidate === button);
        }
        for (const panel of sourcePanels) {
          panel.hidden = panel.dataset.sourcePanel !== button.dataset.sourceTarget;
        }
      });
    }
  </script>
</body>
</html>
""")


def _depths(root: Tensor) -> dict[Tensor, int]:
    depths: dict[Tensor, int] = {}

    def visit(node: Tensor) -> int:
        if node in depths:
            return depths[node]
        depth = 0 if not node._parents else 1 + max(visit(p) for p in node._parents)
        depths[node] = depth
        return depth

    visit(root)
    return depths


def _display_name(node: Tensor, number: int) -> str:
    return node.label or f"v{number}"


def _layout(
    order: list[Tensor], depths: dict[Tensor, int]
) -> tuple[dict[Tensor, tuple[float, float]], int, int]:
    levels: dict[int, list[Tensor]] = {}
    for node in order:
        levels.setdefault(depths[node], []).append(node)

    x_gap = 70
    y_gap = 45
    margin = 50
    max_rows = max(len(nodes) for nodes in levels.values())
    width = max(700, margin * 2 + (max(depths.values()) + 1) * (_BOX_WIDTH + x_gap))
    height = max(360, margin * 2 + max_rows * (_BOX_HEIGHT + y_gap))

    positions: dict[Tensor, tuple[float, float]] = {}
    for depth, nodes in levels.items():
        group_height = len(nodes) * _BOX_HEIGHT + (len(nodes) - 1) * y_gap
        start_y = (height - group_height) / 2
        for row, node in enumerate(nodes):
            positions[node] = (
                margin + depth * (_BOX_WIDTH + x_gap),
                start_y + row * (_BOX_HEIGHT + y_gap),
            )
    return positions, width, height


def _render_edges(
    order: list[Tensor], positions: dict[Tensor, tuple[float, float]]
) -> str:
    edges: list[str] = []
    for child in order:
        child_x, child_y = positions[child]
        for parent in child._parents:
            parent_x, parent_y = positions[parent]
            edges.append(
                '<path class="edge" marker-end="url(#arrow)" '
                f'd="M {parent_x + _BOX_WIDTH} {parent_y + _BOX_HEIGHT / 2} '
                f'L {child_x} {child_y + _BOX_HEIGHT / 2}" />'
            )
    return "".join(edges)


def _render_boxes(
    order: list[Tensor], positions: dict[Tensor, tuple[float, float]]
) -> str:
    numbers = {node: number for number, node in enumerate(order)}
    boxes: list[str] = []
    for node in order:
        x, y = positions[node]
        name = html.escape(_display_name(node, numbers[node]))
        operation = html.escape(node._op or "leaf")
        css_class = "node trainable" if node.requires_grad else "node constant"
        boxes.append(
            f'<g class="{css_class}">'
            f'<rect x="{x}" y="{y}" width="{_BOX_WIDTH}" height="{_BOX_HEIGHT}" rx="10" />'
            f'<text class="name" x="{x + 12}" y="{y + 23}">{name}</text>'
            f'<text x="{x + 12}" y="{y + 46}">data = {node.data:.6g}</text>'
            f'<text x="{x + 12}" y="{y + 66}">grad = {node.grad:.6g}</text>'
            f'<text class="op" x="{x + _BOX_WIDTH - 10}" y="{y + 23}" '
            f'text-anchor="end">{operation}</text>'
            "</g>"
        )
    return "".join(boxes)


def _expression(node: Tensor, numbers: dict[Tensor, int]) -> str:
    if not node._parents:
        return "input"
    parent_names = [_display_name(parent, numbers[parent]) for parent in node._parents]
    if len(parent_names) == 1:
        if node._op.startswith("**"):
            return f"{parent_names[0]} {node._op}"
        return f"{node._op}({parent_names[0]})"
    return f" {node._op} ".join(parent_names)


def _render_walkthrough(order: list[Tensor]) -> str:
    numbers = {node: number for number, node in enumerate(order)}
    rows: list[str] = []
    for step, node in enumerate(order, start=1):
        name = html.escape(_display_name(node, numbers[node]))
        expression = html.escape(_expression(node, numbers))
        rows.append(
            f'<tr><td class="step">{step}</td><td><strong>{name}</strong></td>'
            f"<td><code>{expression}</code></td><td>{node.data:.6g}</td>"
            f"<td>{node.grad:.6g}</td></tr>"
        )
    return "".join(rows)


def _without_leading_header(source: str) -> str:
    lines = source.splitlines(keepends=True)
    if lines and lines[0].startswith("#"):
        lines.pop(0)
        if lines and not lines[0].strip():
            lines.pop(0)
    return "".join(lines)


def _render_source_browser(source_files: Mapping[str, str]) -> str:
    if not source_files:
        return '<p class="empty">No source files were included.</p>'

    buttons: list[str] = []
    panels: list[str] = []
    for index, (name, source) in enumerate(source_files.items()):
        active = " active" if index == 0 else ""
        hidden = "" if index == 0 else " hidden"
        safe_name = html.escape(name)
        safe_source = html.escape(_without_leading_header(source))
        buttons.append(
            f'<button class="file-button{active}" type="button" '
            f'data-source-target="source-{index}">{safe_name}</button>'
        )
        panels.append(
            f'<pre data-source-panel="source-{index}"{hidden}>'
            f"<code>{safe_source}</code></pre>"
        )
    return (
        f'<aside class="file-list">{"".join(buttons)}</aside>'
        f'<div class="source-view">{"".join(panels)}</div>'
    )


def graph_html(root: Tensor, source_files: Mapping[str, str] | None = None) -> str:
    """Return a standalone HTML tutorial for root's autograd graph."""
    order = root._topological_order()
    positions, width, height = _layout(order, _depths(root))
    return _PAGE_TEMPLATE.substitute(
        width=width,
        height=height,
        edges=_render_edges(order, positions),
        boxes=_render_boxes(order, positions),
        walkthrough=_render_walkthrough(order),
        source_browser=_render_source_browser(source_files or {}),
        example_code=html.escape(_EXAMPLE_SNIPPET),
    )


def save_graph_html(
    root: Tensor,
    path: str | Path,
    source_files: Mapping[str, str] | None = None,
) -> Path:
    """Write root's graph tutorial as standalone HTML and return its path."""
    output = Path(path)
    output.write_text(graph_html(root, source_files), encoding="utf-8")
    return output