yobx.torch.export_options#

class yobx.torch.export_options.ExportOptions(strict: bool = False, tracing: bool = False, jit: bool = False, decomposition_table: str | Dict[Any, Callable[[...], Any]] | None = None, strategy: str | None = None, dynamo: bool = False, aten_as_function: bool | Set[Any] | None = None, remove_inplace: bool = True, allow_untyped_output: bool = False, save_ep: Tuple[str, int] | str | None = None, validate_ep: float | bool = False, backed_size_oblivious: bool | str = 'auto', prefer_deferred_runtime_asserts_over_guards: bool = True, fake: bool = False, tracing_module_leaves: Dict[type, Callable[[Module, str], bool]] | None = None)[source]#

Gathers altogether all the options defining the way to export a model into a graph (not onnx).

Parameters:

  • strict – strict export or not, it only applies if torch.export.export() is called

  • decomposition_table – decomposition_table, a string such as 'default' or 'all', or a custom decomposition dict, see get_decomposition_table, it can 'all', 'default' or a decomposition list

  • dynamo – to use torch._dynamo.export instead of torch.export.export()

  • tracing – use symbolic tracing

  • jit – use jit to get a graph then converts it into a fx graph

  • strategy – to overwrite all the previous parameters with just a value

  • remove_inplace – remove inplace nodes

  • aten_as_function – keeps aten function as local function to keep a faithful translation of the fx graph, it can also be a set of function name the export should export as local function such as torch.ops.aten.scaled_dot_product_attention, the default value get_default_aten_as_function returns a default list of functions to keep as function depending on this opset, if no value is specified, this defaults to the whatever the function mentioned above returns

  • allow_untyped_output – allows output with no shape and/or no type

  • save_ep – to save the exported program; if True, will save the graph as text; can be a tuple (str, int) to avoid saving a model bigger than the desired size

  • validate_ep – validates the exported program with the given inputs; by default the tolerance is 1e-5; use a float to change that value

  • backed_size_oblivious – use torch.fx.experimental._config.patch(backed_size_oblivious=True) to allow dynamic dimension equal to 1

  • prefer_deferred_runtime_asserts_over_guards – see torch.export.export()

  • fake – use fake tensors as inputs

  • tracing_module_leaves – this option is used when the module is traced (tracing=True), it specifies which modules should remain a call_module, see yobx.torch.tracing.CustomTracer.

clone(**kwargs) → ExportOptions[source]#

Makes a copy and updates some of the values.

export(mod: Any, args: Tuple[Any, ...] | None, kwargs: Dict[str, Any] | None, tracing_mode: bool, dynamic_shapes: Any, same_signature: bool, input_names: List[str] | None = None, exc: bool = True, verbose: int = 0) → ExportedProgram | GraphModule[source]#

Exports the model into an exported program.

get_decomposition_table() → Dict[Any, Callable[[...], Any]] | None[source]#

Returns the decomposition table.

For decomposition_table='all', returns None because 'all' triggers torch.export.ExportedProgram.run_decompositions() with no arguments (handled by apply_decompositions()).

need_run_decompositions(exported_program) → Tuple[bool, bool][source]#

Final check to see if we need to run decompositions.

post_process_exported_program(exported_program: ExportedProgram, verbose: int = 0, print_exported_program: bool = False) → ExportedProgram[source]#

Run decompositions, remove inplace operations. The graph is modified inplace.

remove_inplace_nodes(graph: Graph, exported_program: ExportedProgram | None = None, verbose: int = 0) → int[source]#

Post-processing to remove inplace nodes.

Parameters:

  • graph – graph to modify

  • exported_program – if available, it is used in the error message to make it easier to trace the code source

  • verbose – verbosity

Returns:

number of inplace nodes removed or -1 if there are any remaining inplace nodes

validate_exported_program(model: Any, exported_program: ExportedProgram, args: Tuple[Any, ...] | None, kwargs: Dict[str, Any] | None, verbose: int = 0)[source]#

Validates the exported program by running the model.

yobx.torch.export_options.apply_decompositions(exported_program: ExportedProgram, decomposition_table: str | Dict[Any, Callable[[...], Any]] | None, backed_size_oblivious: bool | str) → ExportedProgram[source]#

Applies decompositions to an exported program.

Parameters:

  • exported_program – the exported program to decompose

  • decomposition_table – a string ('default' or 'all') or a dict

  • backed_size_oblivious – whether to use backed_size_oblivious=True

Returns:

the decomposed exported program

yobx.torch.export_options.insert_contiguous_between_transpose_and_view(exported_program: ExportedProgram) → ExportedProgram[source]#

Modifies the module inplace to insert a contiguous node between a transpose node followed by a view node. The modification takes place inplace. See pytorch/pytorch#136543.