xref: /aosp_15_r20/external/pytorch/torch/_dynamo/backends/onnxrt.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# mypy: ignore-errors
2
3# This backend is maintained by ONNX team. To direct issues
4# to the right people, please tag related GitHub issues with `module: onnx`.
5#
6# Maintainers' Github IDs: wschin, xadupre
7from torch.onnx._internal.onnxruntime import (
8    is_onnxrt_backend_supported,
9    torch_compile_backend,
10)
11
12from .registry import register_backend
13
14
15def has_onnxruntime():
16    # FIXME: update test/dynamo/test_backends.py to call is_onnxrt_backend_supported()
17    return is_onnxrt_backend_supported()
18
19
20if is_onnxrt_backend_supported():
21    register_backend(name="onnxrt", compiler_fn=torch_compile_backend)
22else:
23
24    def information_displaying_backend(*args, **kwargs):
25        raise ImportError(
26            "onnxrt is not registered as a backend. "
27            "Please make sure all dependencies such as "
28            "numpy, onnx, onnxscript, and onnxruntime-training are installed. "
29            "Suggested procedure to fix dependency problem:\n"
30            "  (1) pip or conda install numpy onnx onnxscript onnxruntime-training.\n"
31            "  (2) Open a new python terminal.\n"
32            "  (3) Call the API `torch.onnx.is_onnxrt_backend_supported()`:\n"
33            "  (4)   If it returns `True`, then you can use `onnxrt` backend.\n"
34            "  (5)   If it returns `False`, please execute the package importing section in "
35            "torch/onnx/_internal/onnxruntime.py under pdb line-by-line to see which import fails."
36        )
37
38    register_backend(name="onnxrt", compiler_fn=information_displaying_backend)
39