checkpoint(function: Callable[..., 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 single
lucid.Tensor.Positional tensor inputs to function.
preserve_rng_statebool= TrueAccepted for API compatibility. RNG state restoration is not yet
implemented — set to
False when function contains stochastic
layers.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
TensorOutput of function(*args, **kwargs).
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.]])