pytorch
865d30a3 - [ONNX] Add new ops and enable the MNIST test (#99284)

Commit
1 year ago
[ONNX] Add new ops and enable the MNIST test (#99284) <!-- copilot:all --> ### <samp>πŸ€– Generated by Copilot at d490881</samp> ### Summary πŸ“¦πŸš€πŸ—ΊοΈ <!-- 1. πŸ“¦ for the package name change and update. 2. πŸš€ for the test case enablement and feature support. 3. πŸ—ΊοΈ for the operator mappings addition. --> This pull request adds support for converting more PyTorch FX operators to ONNX using the `onnxscript` package. It updates the package installation script and the operator mappings in `function_dispatcher.py`. It also enables a test case for the `max_pool2d` operator. > _We unleash the fury of the ONNX script_ > _We break the chains of the skipped test_ > _We map the functions of the ATen torch_ > _We forge the metal of the ONNX forge_ ### Walkthrough * Fix and update the installation of the `onnxscript` package ([link](https://github.com/pytorch/pytorch/pull/99284/files?diff=unified&w=0#diff-ddc2d714e91795d6bfc059caca91ab22d0740909f744d4c22f6642a71b09e8ecL29-R30)) * Add more mappings from PyTorch ATen operators to ONNX operators in the `_ATENLIB_FUNCTIONS` dictionary in the `function_dispatcher.py` module ([link](https://github.com/pytorch/pytorch/pull/99284/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL44-R91), [link](https://github.com/pytorch/pytorch/pull/99284/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL78-R130), [link](https://github.com/pytorch/pytorch/pull/99284/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL103-R193), [link](https://github.com/pytorch/pytorch/pull/99284/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL131-R210), [link](https://github.com/pytorch/pytorch/pull/99284/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL144-R223)) * Enable the `test_max_pool2d` method in the `TestFxToOnnx` class in the `test_fx_to_onnx.py` file, since the `max_pool2d` operator is now supported by the `onnxscript` package ([link](https://github.com/pytorch/pytorch/pull/99284/files?diff=unified&w=0#diff-9dde99e3ce414c7ca10dd319e22f74bcaaddaccd3ab6560c24777baf20616d27L71-L73)) Pull Request resolved: https://github.com/pytorch/pytorch/pull/99284 Approved by: https://github.com/BowenBao
Author
Committer
Parents
Loading