pytorch_pfn_extras.ops.register#

pytorch_pfn_extras.ops.register(name, fwd_op, bwd_op)#

Register a custom op under torch.ops.ppe.name

The function appears as a primitive op in the forward and backward torch.fx.Graph``s after compiling torch code with `aot_autograd` backend. Note that for backward functions, all the arguments of the backward pass together with the forward arguments are passed to it. This means if forward had ``fwd_op(x, y) x,y arguments, the custom bwd_op needs to have a signature like``bwd_op(grad_output, x, y)``

Parameters:

  • name (str) – name of the op, shows how it is registered in torch.ops.ppe.

  • fwd_op (ppe.ops.OpDesc) – code that is executed in the forward pass

  • bwd_op (ppe.ops.OpDesc) – code that is executed in the backward pass

Return type:

None