checkpoint(function: Callable[..., Tensor | tuple[Tensor, ...]], args: Tensor = (), preserve_rng_state: bool = True, use_reentrant: bool = True, kwargs: object = {})Run function under gradient checkpointing.
Executes function(*args, **kwargs) during the forward pass
without tracking intermediate activations (no_grad context).
During the backward pass the function is re-executed under
enable_grad to reconstruct the local autograd graph, and the
gradients are computed through that graph.
Parameters
functioncallableThe differentiable segment to checkpoint. Must accept tensors as
positional arguments and return a
lucid.Tensor or a
tuple of them.Positional tensor inputs to function.
preserve_rng_statebool= TrueRestore the random state the forward saw before recomputing, so
stochastic layers draw the same values twice. Default
True.use_reentrantbool= TrueTrue (default, kept for backward compatibility) ties the
recomputation to the positional inputs: if none of them requires
grad, backward never runs and nothing inside the segment receives
a gradient — including parameters it closed over. False
anchors the node in the graph so the segment is always recomputed
and those parameters accumulate normally. Prefer False
unless something depends on the old behaviour.**kwargsobject= {}Extra keyword arguments forwarded to function on both the
forward and recomputation passes.
Returns
Examples
>>> import lucid
>>> W = lucid.tensor([[1.0, 0.0], [0.0, 1.0]], requires_grad=True)
>>> b = lucid.tensor([0.0, 0.0], requires_grad=True)
>>> x = lucid.tensor([[1.0, -2.0]], requires_grad=True)
>>> def block(x):
... return lucid.nn.functional.relu(x @ W + b)
>>> y = lucid.autograd.checkpoint(block, x, use_reentrant=False)
>>> y.sum().backward()
>>> x.grad
tensor([[1., 0.]])