torch.registerUnaryOp
function registerUnaryOp<Op extends UnaryOpNames>(config: UnaryOpConfig<Op>): UnaryOpFnRegister a unary operation with minimal boilerplate.
Returns the registered function so it can be exported from the module.
Parameters
configUnaryOpConfig<Op>
Returns
Examples
export const exp = registerUnaryOp({
op: 'exp',
cpu: Math.exp,
shaderEntry: 'exp_op',
backward: (ctx, grad) => ({
input: grad.mul(ctx.saved_tensors[0]),
}),
saveForBackward: 'output',
});