warp.func_replay

warp.func_replay(forward_fn)

Decorator to register a custom replay function for a given forward function. The replay function is the function version that is called in the forward phase of the backward pass (replay mode) and corresponds to the forward function by default. The provided function has to match the signature of one of the original forward function overloads.