torch.compiler.substitute_in_graph#

torch.compiler.substitute_in_graph(original_fn,*,can_constant_fold_through=False,skip_signature_check=False)[source]#

Register a polyfill handler for a function, usually a C function from the C extension, to beused in place of the original function when inlining the original function in the graph.

Note

The polyfill handler is only used when inlining the original function. It is not used whenthe original function is called directly. In the eager mode, the decorated function callsthe performant C function rather than the polyfill handler.

The polyfill handler is a function that will be called in place of the original function wheninlining the original function. The polyfill handler should have the same signature and the samebehavior as the original function.

Parameters

• original_fn (callable) – The original function, usually a C function, to register a polyfillhandler for.

• can_constant_fold_through (bool,optional) – Whether the polyfill handler can be constantfolded through. That is, if the polyfill handler is a pure function and its argumentsare constant, the result of the polyfill handler can be constant folded during thecompilation. Defaults toFalse.

• skip_signature_check (bool,optional) – Whether to skip the signature check between theoriginal function and the polyfill handler. Defaults toFalse.

Returns

A decorator that registers the polyfill handler for the original function.

Return type

Callable[[Callable[[~_P],_R]],Callable[[~_P],_R]]

Example:

>>>importoperator>>>operator.indexOf([1,2,3,4,5],3)2>>>torch.compile(operator.indexOf,fullgraph=True)([1,2,3,4,5],3)...# xdoctest: +SKIP("Long tracebacks")Traceback (most recent call last):...torch._dynamo.exc.Unsupported:...>>>@torch.compiler.substitute_in_graph(operator.indexOf)...defindexOf(a,b,/):...fori,iteminenumerate(a):...ifitemisboritem==b:...returni...raiseValueError("sequence.index(x): x not in sequence")>>>>>>torch.compile(operator.indexOf,fullgraph=True)([1,2,3,4,5],3)2