torch.register_backward
function register_backward<Op extends OpName>(config: AutogradConfig<Op>): AutogradHandleImpl<Op>function register_backward<Op extends OpName>(op: Op, setup_context: (ctx: GradContext<Op>, inputs: InputsFor<Op>, output: TensorLike) => void, backward: (ctx: GradContext<Op>, grad_output: TensorLike) => GradientsFor<Op>, options: RegisterBackwardOptions): AutogradHandleImpl<Op>Register autograd with full configuration object.
Parameters
configAutogradConfig<Op>
Returns
Examples
torch.library.register_backward({
op: 'mul',
setup_context: (ctx, inputs, output) => {
ctx.save_for_backward(inputs.a, inputs.b);
},
backward: (ctx, grad_output) => {
const [a, b] = ctx.saved_tensors;
return {
a: torch.mul(grad_output, b),
b: torch.mul(grad_output, a),
};
},
});