thunder.torch._register_custom_op

thunder.torch._register_custom_op(custom_op)

Register custom_op()’ed or triton_op()’ed function to Thunder.

torch.library.custom_op() operators always have schema as per https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit?tab=t.0#heading=h.aotf6sastysc saying “Use torch.library.custom_op to decorate a function to turn it into a custom operator. The function must be decorated with type annotations, and you must correctly annotate inputs that are being mutated.” An example schema is my_lib::foo(Tensor a, Tensor b, float? c=None, str d="") -> Tensor.

This function does three things:

1. Register custom_op as a new Symbol whose fn is custom_op._opoverload and meta is based on custom_op._abstract_fn

2. When _setup_context_fn and _backward_fn are defined, register augmented forward and backward through thunder.core.transforms.register_augmented_forward() and thunder.core.transforms.register_backward().

Parameters:

• custom_op (CustomOpDef) – torch.library.custom_op()’ed function. This is not torch.ops.{namespace}.{name}.

• custom_op (CustomOpDef) –

Returns:

A symbol representing the input custom_op.

Return type:

Symbol

Note

This feature is experimental and subject to change.