xnn.common.models.recompute.recompute#

xnn.common.models.recompute.recompute(fn, *tensors, grad_fn=None)[source]#

Evaluate fn(*tensors) as a block that is recomputed rather than stored.

Parameters:

fn (callable) – A function of tensors only (bind other arguments with a closure or functools.partial()) returning one tensor.
*tensors (Tensor) – Its inputs; gradients flow to those with requires_grad.
grad_fn (callable, optional) – Analytic gradient grad_fn(grad_out, *tensors) -> tuple with one entry per input (None for inputs without gradient), written in differentiable torch ops. Used in place of autograd through fn for the first derivative; the second derivative is autograd over it.

Returns:

fn(*tensors), with first and second derivatives evaluated by re-running the block, so the retained memory is its inputs only.

Return type:

Tensor