onnx_diagnostic.torch_export_patches.patch_module¶

class onnx_diagnostic.torch_export_patches.patch_module.RewriteControlFlow(wrapper_name)[source][source]¶

generic_visit(node)[source][source]¶: Called if no explicit visitor function exists for a node.

onnx_diagnostic.torch_export_patches.patch_module.transform_method(func, wrapper_name='torch_cond')[source][source]¶

Returns a new function based on func where every test (if, while, assert, ternary, comparison, boolean op) is replaced by a call to wrapper_name.

wrapper_name should refer to a function taking a single boolean argument.