A reusable forward op with a hand-written backward.
fwd(inputs) returns the forward outputs. vjp(primals, cotangents, outputs) returns one gradient per primal (in primal order). Both must be
pure functions of their arguments — the vjp re-reads everything from the
graph during the enclosing backward.
A reusable forward op with a hand-written backward.
fwd(inputs)returns the forward outputs.vjp(primals, cotangents, outputs)returns one gradient per primal (in primal order). Both must be pure functions of their arguments — the vjp re-reads everything from the graph during the enclosing backward.