yobx.torch.interpreter.dispatcher#

class yobx.torch.interpreter.dispatcher.Dispatcher(registered_functions: Dict[type | str, Callable], verbose: int = 0)[source]#

Used to changes the way class DynamoInterpreter selects the function translating aten function or module.

Parameters:

  • registered_functions – registered functions

  • verbose – verbose

fallback(name: Any, fct: Callable | None, args: List[Any], kwargs: Dict[str, Any], builder: GraphBuilder) → Callable | None[source]#

The function is called after the function converting an aten function into ONNX. fct is this function. It can be changed and just set when mapping was found.

Parameters:

  • name – object or str

  • fct – function found so far

  • args – known arguments coming from the graph module

  • kwargs – known named arguments coming from the graph module

  • builder – GraphBuilder

Returns:

callable

find_function(name: Any) → Callable | None[source]#

Finds the most suitable function to translate a function.

Parameters:

name – function name or definition

Returns:

the function or None if not found

The signature of the returned function is similar to a function such as aten_elu.

find_method(name: Any) → Callable | None[source]#

Finds the most suitable function to translate a method.

Parameters:

name – method name or definition

Returns:

the function or None if not found

The signature of the returned function is similar to a function such as aten_elu.

merge(other_dispatcher: Dispatcher)[source]#

Adds other dispateched functions into this one.

property supported: Set[type | str]#

Returns the list supported dispateched names.

class yobx.torch.interpreter.dispatcher.ForceDispatcher(signatures: Dict[str, Callable] | None = None, verbose: int = 0, domain: str = 'aten.lib', version: int = 1, strict: bool = False, only_registered: bool = False)[source]#

Implements a dispatcher which fails whenever there is no converting for a node in the fx graph. There is no fallback to the existing functions. When no function is found, an onnx node is added with a non standard domain.

Parameters:

  • signatures – function used only for their signature mapping a name to a function in order to have parameter names

  • verbose – verbose

  • domain – domain of the added node

  • version – version of the domain

  • strict – when an input is not a tensor, it becomes a named parameter if strict is False

  • only_registered – fails if a function is not found in signatures

fallback(name: Any, fct: Callable | None, args: List[Any], kwargs: Dict[str, Any], builder: GraphBuilder) → Callable | None[source]#

The function is called after the function converting an aten function into ONNX. fct is this function. It can be changed and just set when mapping was found.

Parameters:

  • name – object or str

  • fct – function found so far

  • args – known arguments coming from the graph module

  • kwargs – known named arguments coming from the graph module

  • builder – GraphBuilder

Returns:

callable