python/dpt/autodiff.py

1"""First-order tape ownership for composed deterministic operators.23A tape is an execution lifetime, not an image Jacobian. This small owner makes4output seeding, error checkpoints and discard/reset explicit. Scientific5operators still receive the underlying tape through their tape= argument.6"""78from __future__ import annotations910from collections.abc import Sequence11from dataclasses import dataclass, field12from typing import Any, Protocol1314from dpt._runtime import load_warp, prepare_context, require_no_tape15from dpt.contracts import ContractError161718class CheckedWorkspace(Protocol):19    def clear_status(self) -> None: ...20    def check_status(self) -> None: ...212223# region book:autodiff-first-order-lifetime24@dataclass(slots=True)25class FirstOrderPass:26    """Own one forward/reverse evaluation and its explicit diagnostic checkpoint.2728    Use as a context manager; pass .tape into canonical forward operators.29    Call backward after exiting the context. Keep original inputs immutable30    until close(). A new iteration requires a fresh context entry after close.31    """3233    workspaces: Sequence[CheckedWorkspace]34    tape: Any = field(default=None, init=False, repr=False)35    _state: str = field(default="new", init=False)3637    def __enter__(self) -> FirstOrderPass:38        if self._state != "new":39            raise ContractError("a FirstOrderPass owns exactly one evaluation")40        require_no_tape()41        for workspace in self.workspaces:42            workspace.clear_status()43        self.tape = load_warp().Tape()44        self.tape.__enter__()45        self._state = "recording"46        return self4748    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:49        self.tape.__exit__(exc_type, exc_value, traceback)50        self._state = "recorded"51        if exc_type is not None:52            self.close()5354    def backward(self, output: Any, seed: Any) -> None:55        """Seed a caller-owned output cotangent and complete numerical checks."""56        if self._state != "recorded":57            raise ContractError("backward requires one completed, unconsumed forward pass")58        try:59            wp = load_warp()60            if not isinstance(output, wp.array) or output.grad is None:61                raise ContractError("output must be a Warp CUDA array with a preallocated gradient")62            ctx = prepare_context(device=output.device)63            ctx.array(output, "output", dtype=output.dtype, ndim=output.ndim)64            ctx.array(seed, "seed", dtype=output.dtype, shape=tuple(output.shape), ndim=output.ndim)65            self.tape.backward(grads={output: seed})66            for workspace in self.workspaces:67                workspace.check_status()68            self._state = "complete"69        except BaseException:70            self.close()71            raise7273    def close(self) -> None:74        """Clear accumulated cotangents and release every abandoned recording."""75        if self._state == "recording":76            raise ContractError("exit the recording context before closing it")77        if self.tape is not None and self._state != "closed":78            self.tape.zero()79            self.tape.reset()80            for workspace in self.workspaces:81                discard = getattr(workspace, "discard_recording", None)82                if discard is not None:83                    discard(self.tape)84        self._state = "closed"858687# endregion book:autodiff-first-order-lifetime88