onnx_diagnostic.torch_export_patches.patches.patch_torch

onnx_diagnostic.torch_export_patches.patches.patch_torch.patch__check_input_constraints_for_graph(previous_function: Callable, input_placeholders: list[Node], flat_args_with_path, range_constraints, verbose: int = 0) → None

class onnx_diagnostic.torch_export_patches.patches.patch_torch.patched_ShapeEnv

onnx_diagnostic.torch_export_patches.patches.patch_torch.patched__broadcast_in_dim_meta(a: Tensor, shape: Size | list[int] | tuple[int, ...], broadcast_dimensions: Sequence[int])

Patches torch._prims._broadcast_in_dim_meta.

onnx_diagnostic.torch_export_patches.patches.patch_torch.patched__broadcast_shapes(*_shapes)

Patches torch._refs._broadcast_shapes.

onnx_diagnostic.torch_export_patches.patches.patch_torch.patched__constrain_user_specified_dimhint_range(symint: SymInt, hint: int, dim: _DimHint, range_constraints, shape_env, keypath: KeyPath, i: int | None = None) → str | None

Patches torch._export.non_strict_utils._constrain_user_specified_dimhint_range.

onnx_diagnostic.torch_export_patches.patches.patch_torch.patched__maybe_broadcast(*args, preserve_cpu_scalar_tensors=True)

Patches torch._refs._maybe_broadcast.

onnx_diagnostic.torch_export_patches.patches.patch_torch.patched_infer_size(a, b)

Patches torch._subclasses.fake_impls.infer_size.

onnx_diagnostic.torch_export_patches.patches.patch_torch.patched_vmap(func, in_dims=0, out_dims=0)

Python implementation of torch.vmap(). The implementation raises an issue when it is being exported with torch.export.export() when the function is called with non tensors arguments and the batch size is dynamic.

onnx_diagnostic.torch_export_patches.patches.patch_torch.retrieve_stacktrace()

Retrieves and prints the current stack trace, avoids every torch file.