onnxscript 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (184) hide show
  1. onnxscript-0.1.0/LICENSE +21 -0
  2. onnxscript-0.1.0/MANIFEST.in +11 -0
  3. onnxscript-0.1.0/PKG-INFO +370 -0
  4. onnxscript-0.1.0/README.md +318 -0
  5. onnxscript-0.1.0/VERSION +1 -0
  6. onnxscript-0.1.0/onnxscript/__init__.py +131 -0
  7. onnxscript-0.1.0/onnxscript/_framework_apis/__init__.py +3 -0
  8. onnxscript-0.1.0/onnxscript/_framework_apis/torch_2_5.py +117 -0
  9. onnxscript-0.1.0/onnxscript/_framework_apis/torch_2_6.py +45 -0
  10. onnxscript-0.1.0/onnxscript/_internal/__init__.py +0 -0
  11. onnxscript-0.1.0/onnxscript/_internal/analysis.py +229 -0
  12. onnxscript-0.1.0/onnxscript/_internal/ast_utils.py +64 -0
  13. onnxscript-0.1.0/onnxscript/_internal/autocast.py +250 -0
  14. onnxscript-0.1.0/onnxscript/_internal/deprecation.py +78 -0
  15. onnxscript-0.1.0/onnxscript/_internal/param_manipulation.py +148 -0
  16. onnxscript-0.1.0/onnxscript/_internal/runtime_typing.py +43 -0
  17. onnxscript-0.1.0/onnxscript/_internal/utils.py +99 -0
  18. onnxscript-0.1.0/onnxscript/_internal/version_utils.py +118 -0
  19. onnxscript-0.1.0/onnxscript/_legacy_ir/__init__.py +341 -0
  20. onnxscript-0.1.0/onnxscript/_legacy_ir/visitor.py +937 -0
  21. onnxscript-0.1.0/onnxscript/_thirdparty/asciichartpy.py +313 -0
  22. onnxscript-0.1.0/onnxscript/backend/__init__.py +2 -0
  23. onnxscript-0.1.0/onnxscript/backend/onnx_backend.py +303 -0
  24. onnxscript-0.1.0/onnxscript/backend/onnx_export.py +885 -0
  25. onnxscript-0.1.0/onnxscript/converter.py +1470 -0
  26. onnxscript-0.1.0/onnxscript/evaluator.py +619 -0
  27. onnxscript-0.1.0/onnxscript/function_libs/tools/torch_lib/deduce_type_constraints.py +403 -0
  28. onnxscript-0.1.0/onnxscript/function_libs/tools/torch_lib/generate_aten_signatures.py +333 -0
  29. onnxscript-0.1.0/onnxscript/function_libs/tools/torch_lib/generate_prims_signatures.py +331 -0
  30. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/__init__.py +12 -0
  31. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/_constants.py +5 -0
  32. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/_flags.py +58 -0
  33. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/graph_building/__init__.py +56 -0
  34. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/graph_building/_graph_building_ir.py +723 -0
  35. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/graph_building/_graph_building_torch.py +1125 -0
  36. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/__init__.py +27 -0
  37. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/common.py +80 -0
  38. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/core.py +8935 -0
  39. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/fft.py +385 -0
  40. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/linalg.py +399 -0
  41. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/nested.py +25 -0
  42. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/nn.py +2713 -0
  43. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/prims.py +850 -0
  44. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/quantized_decomposed.py +63 -0
  45. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/sparse.py +23 -0
  46. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/special.py +387 -0
  47. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/ops/vision.py +25 -0
  48. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/registration.py +151 -0
  49. onnxscript-0.1.0/onnxscript/function_libs/torch_lib/tensor_typing.py +74 -0
  50. onnxscript-0.1.0/onnxscript/ir/__init__.py +153 -0
  51. onnxscript-0.1.0/onnxscript/ir/_convenience.py +447 -0
  52. onnxscript-0.1.0/onnxscript/ir/_core.py +3119 -0
  53. onnxscript-0.1.0/onnxscript/ir/_display.py +49 -0
  54. onnxscript-0.1.0/onnxscript/ir/_enums.py +163 -0
  55. onnxscript-0.1.0/onnxscript/ir/_graph_comparison.py +23 -0
  56. onnxscript-0.1.0/onnxscript/ir/_io.py +97 -0
  57. onnxscript-0.1.0/onnxscript/ir/_linked_list.py +276 -0
  58. onnxscript-0.1.0/onnxscript/ir/_metadata.py +44 -0
  59. onnxscript-0.1.0/onnxscript/ir/_name_authority.py +72 -0
  60. onnxscript-0.1.0/onnxscript/ir/_polyfill.py +25 -0
  61. onnxscript-0.1.0/onnxscript/ir/_protocols.py +598 -0
  62. onnxscript-0.1.0/onnxscript/ir/_schemas.py +548 -0
  63. onnxscript-0.1.0/onnxscript/ir/_tape.py +120 -0
  64. onnxscript-0.1.0/onnxscript/ir/_type_casting.py +106 -0
  65. onnxscript-0.1.0/onnxscript/ir/convenience.py +32 -0
  66. onnxscript-0.1.0/onnxscript/ir/external_data.py +396 -0
  67. onnxscript-0.1.0/onnxscript/ir/passes/__init__.py +33 -0
  68. onnxscript-0.1.0/onnxscript/ir/passes/_pass_infra.py +172 -0
  69. onnxscript-0.1.0/onnxscript/ir/serde.py +1620 -0
  70. onnxscript-0.1.0/onnxscript/ir/tensor_adapters.py +122 -0
  71. onnxscript-0.1.0/onnxscript/ir/traversal.py +82 -0
  72. onnxscript-0.1.0/onnxscript/irbuilder.py +542 -0
  73. onnxscript-0.1.0/onnxscript/main.py +167 -0
  74. onnxscript-0.1.0/onnxscript/onnx_opset/__init__.py +232 -0
  75. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset1.py +4100 -0
  76. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset10.py +1227 -0
  77. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset11.py +4013 -0
  78. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset12.py +1078 -0
  79. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset13.py +3924 -0
  80. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset14.py +999 -0
  81. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset15.py +604 -0
  82. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset16.py +1255 -0
  83. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset17.py +561 -0
  84. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset18.py +1803 -0
  85. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset19.py +1942 -0
  86. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset2.py +218 -0
  87. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset20.py +675 -0
  88. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset21.py +1976 -0
  89. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset22.py +2588 -0
  90. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset3.py +199 -0
  91. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset4.py +77 -0
  92. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset5.py +84 -0
  93. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset6.py +944 -0
  94. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset7.py +1243 -0
  95. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset8.py +444 -0
  96. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset9.py +1485 -0
  97. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml1.py +974 -0
  98. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml2.py +112 -0
  99. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml3.py +308 -0
  100. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml4.py +129 -0
  101. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml5.py +158 -0
  102. onnxscript-0.1.0/onnxscript/onnx_opset/_impl/opset_ai_onnx_preview_training1.py +577 -0
  103. onnxscript-0.1.0/onnxscript/onnx_types.py +229 -0
  104. onnxscript-0.1.0/onnxscript/optimizer/__init__.py +39 -0
  105. onnxscript-0.1.0/onnxscript/optimizer/_constant_folding.py +1083 -0
  106. onnxscript-0.1.0/onnxscript/optimizer/_inliner.py +312 -0
  107. onnxscript-0.1.0/onnxscript/optimizer/_legacy/_optimizer.py +98 -0
  108. onnxscript-0.1.0/onnxscript/optimizer/_legacy/_remove_unused_proto.py +144 -0
  109. onnxscript-0.1.0/onnxscript/optimizer/_legacy/_simple_function_folding.py +243 -0
  110. onnxscript-0.1.0/onnxscript/optimizer/_legacy/constant_folding.py +293 -0
  111. onnxscript-0.1.0/onnxscript/optimizer/_legacy/evaluator.py +439 -0
  112. onnxscript-0.1.0/onnxscript/optimizer/_optimizer.py +61 -0
  113. onnxscript-0.1.0/onnxscript/optimizer/_remove_unused.py +106 -0
  114. onnxscript-0.1.0/onnxscript/optimizer/_remove_unused_function.py +72 -0
  115. onnxscript-0.1.0/onnxscript/py.typed +1 -0
  116. onnxscript-0.1.0/onnxscript/rewriter/__init__.py +56 -0
  117. onnxscript-0.1.0/onnxscript/rewriter/_ir_utils.py +111 -0
  118. onnxscript-0.1.0/onnxscript/rewriter/broadcast_to_matmul.py +180 -0
  119. onnxscript-0.1.0/onnxscript/rewriter/cast_constant_of_shape.py +50 -0
  120. onnxscript-0.1.0/onnxscript/rewriter/collapse_slices.py +141 -0
  121. onnxscript-0.1.0/onnxscript/rewriter/erfgelu.py +27 -0
  122. onnxscript-0.1.0/onnxscript/rewriter/function_rule.py +232 -0
  123. onnxscript-0.1.0/onnxscript/rewriter/gemm_to_matmul_add.py +21 -0
  124. onnxscript-0.1.0/onnxscript/rewriter/generic_pattern.py +700 -0
  125. onnxscript-0.1.0/onnxscript/rewriter/llama_rule_sets.py +287 -0
  126. onnxscript-0.1.0/onnxscript/rewriter/no_op.py +56 -0
  127. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/__init__.py +50 -0
  128. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/bfloat16_utils/bfloat16_converter.py +99 -0
  129. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/fused_matmul_rule_sets.py +177 -0
  130. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/group_normalization_merge_silu.py +64 -0
  131. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/instance_to_group_normalization.py +155 -0
  132. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/softmax.py +63 -0
  133. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/transformers/__init__.py +21 -0
  134. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/transformers/biassplitgelu.py +31 -0
  135. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/transformers/fastgelu.py +29 -0
  136. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/transformers/layernorm.py +47 -0
  137. onnxscript-0.1.0/onnxscript/rewriter/onnxruntime/transformers/multihead_attention.py +715 -0
  138. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/__init__.py +9 -0
  139. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/_core.py +28 -0
  140. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/_smollm_1.py +253 -0
  141. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/_smollm_2.py +467 -0
  142. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/_test_models.py +122 -0
  143. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/_test_utils.py +42 -0
  144. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/cos_sin_cache.py +154 -0
  145. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/gqa.py +156 -0
  146. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/mha.py +198 -0
  147. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/rms_normalization.py +95 -0
  148. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/rotary_embedding.py +64 -0
  149. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/sdpa.py +75 -0
  150. onnxscript-0.1.0/onnxscript/rewriter/ort_fusions/skip_normalization.py +46 -0
  151. onnxscript-0.1.0/onnxscript/rewriter/pattern.py +1714 -0
  152. onnxscript-0.1.0/onnxscript/rewriter/testing.py +77 -0
  153. onnxscript-0.1.0/onnxscript/sourceinfo.py +59 -0
  154. onnxscript-0.1.0/onnxscript/tensor.py +227 -0
  155. onnxscript-0.1.0/onnxscript/testing/__init__.py +482 -0
  156. onnxscript-0.1.0/onnxscript/tools/__init__.py +4 -0
  157. onnxscript-0.1.0/onnxscript/tools/benchmark/__init__.py +23 -0
  158. onnxscript-0.1.0/onnxscript/tools/benchmark/benchmark_helpers.py +783 -0
  159. onnxscript-0.1.0/onnxscript/tools/benchmark/benchmark_run.py +140 -0
  160. onnxscript-0.1.0/onnxscript/tools/benchmark/export_model.py +207 -0
  161. onnxscript-0.1.0/onnxscript/tools/benchmark/export_model_batch.py +146 -0
  162. onnxscript-0.1.0/onnxscript/tools/memory_peak.py +244 -0
  163. onnxscript-0.1.0/onnxscript/tools/training_helper.py +50 -0
  164. onnxscript-0.1.0/onnxscript/tools/transformers_models/__init__.py +190 -0
  165. onnxscript-0.1.0/onnxscript/tools/transformers_models/llama.py +168 -0
  166. onnxscript-0.1.0/onnxscript/tools/transformers_models/mistral.py +238 -0
  167. onnxscript-0.1.0/onnxscript/tools/transformers_models/phi.py +248 -0
  168. onnxscript-0.1.0/onnxscript/tools/transformers_models/phi3.py +259 -0
  169. onnxscript-0.1.0/onnxscript/type_annotation.py +281 -0
  170. onnxscript-0.1.0/onnxscript/utils/__init__.py +0 -0
  171. onnxscript-0.1.0/onnxscript/utils/evaluation_utils.py +56 -0
  172. onnxscript-0.1.0/onnxscript/utils/timing_utils.py +33 -0
  173. onnxscript-0.1.0/onnxscript/utils/utils.py +84 -0
  174. onnxscript-0.1.0/onnxscript/values.py +790 -0
  175. onnxscript-0.1.0/onnxscript/version_converter/__init__.py +21 -0
  176. onnxscript-0.1.0/onnxscript/version_converter/_version_converter.py +314 -0
  177. onnxscript-0.1.0/onnxscript.egg-info/PKG-INFO +370 -0
  178. onnxscript-0.1.0/onnxscript.egg-info/SOURCES.txt +182 -0
  179. onnxscript-0.1.0/onnxscript.egg-info/dependency_links.txt +1 -0
  180. onnxscript-0.1.0/onnxscript.egg-info/requires.txt +5 -0
  181. onnxscript-0.1.0/onnxscript.egg-info/top_level.txt +1 -0
  182. onnxscript-0.1.0/pyproject.toml +227 -0
  183. onnxscript-0.1.0/setup.cfg +4 -0
  184. onnxscript-0.1.0/setup.py +34 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) Microsoft Corporation
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,11 @@
1
+ # include
2
+ include *.rst
3
+ include LICENSE
4
+ include README.md
5
+ include VERSION
6
+ include MANIFEST.in
7
+
8
+ # exclude from sdist
9
+ recursive-exclude onnxscript *.onnx
10
+ recursive-exclude onnxscript *_test.py
11
+ prune */__pycache__
@@ -0,0 +1,370 @@
1
+ Metadata-Version: 2.2
2
+ Name: onnxscript
3
+ Version: 0.1.0
4
+ Summary: Naturally author ONNX functions and models using a subset of Python
5
+ Author-email: Microsoft Corporation <onnx@microsoft.com>
6
+ License: MIT License
7
+
8
+ Copyright (c) Microsoft Corporation
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+
28
+ Project-URL: Homepage, https://onnxscript.ai/
29
+ Project-URL: Repository, https://github.com/microsoft/onnxscript
30
+ Classifier: Development Status :: 4 - Beta
31
+ Classifier: Environment :: Console
32
+ Classifier: Intended Audience :: Developers
33
+ Classifier: Operating System :: POSIX
34
+ Classifier: Operating System :: MacOS :: MacOS X
35
+ Classifier: Operating System :: Microsoft :: Windows
36
+ Classifier: Programming Language :: Python :: 3.8
37
+ Classifier: Programming Language :: Python :: 3.9
38
+ Classifier: Programming Language :: Python :: 3.10
39
+ Classifier: Programming Language :: Python :: 3.11
40
+ Classifier: Programming Language :: Python :: 3.12
41
+ Classifier: Programming Language :: Python :: 3.13
42
+ Classifier: License :: OSI Approved :: MIT License
43
+ Requires-Python: >=3.8
44
+ Description-Content-Type: text/markdown
45
+ License-File: LICENSE
46
+ Requires-Dist: numpy
47
+ Requires-Dist: onnx>=1.16
48
+ Requires-Dist: typing_extensions>=4.10
49
+ Requires-Dist: ml_dtypes
50
+ Requires-Dist: packaging
51
+ Dynamic: project-url
52
+
53
+ # ONNX Script
54
+
55
+ [![CI](https://github.com/microsoft/onnxscript/actions/workflows/main.yaml/badge.svg)](https://github.com/microsoft/onnxscript/actions/workflows/main.yaml)
56
+ [![Dev Release](https://aiinfra.visualstudio.com/ONNX%20Converters/_apis/build/status%2Fonnxscript-release-dev?branchName=main&label=Dev%20Release)](https://aiinfra.visualstudio.com/ONNX%20Converters/_build/latest?definitionId=1258&branchName=main)
57
+ [![PyPI - Version](https://img.shields.io/pypi/v/onnxscript.svg)](https://pypi.org/project/onnxscript)
58
+ [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/onnxscript.svg)](https://pypi.org/project/onnxscript)
59
+ [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
60
+ [![Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
61
+
62
+ ONNX Script enables developers to naturally author ONNX functions and
63
+ models using a subset of Python. ONNX Script is:
64
+
65
+ * **Expressive:** enables the authoring of all ONNX functions.
66
+ * **Simple and concise:** function code is natural and simple.
67
+ * **Debuggable:** allows for eager-mode evaluation that provides for a
68
+ more delightful ONNX model debugging experience.
69
+
70
+ This repo also covers:
71
+
72
+ * **ONNX IR:** an in-memory IR that supports the full ONNX spec, designed
73
+ for graph construction, analysis and transformation.
74
+ * **ONNX Script Optimizer:** provides functionality to optimize an ONNX
75
+ model by performing optimizations and clean-ups such as constant folding,
76
+ dead code elimination, etc.
77
+ * **ONNX Rewriter:** provides functionality to replace certain patterns in
78
+ an ONNX graph with replacement patterns based on user-defined rewrite rules.
79
+
80
+ Note however that ONNX Script does **not** intend to support the entirety
81
+ of the Python language.
82
+
83
+ Website: [https://onnxscript.ai/](https://onnxscript.ai/)
84
+
85
+ ## Design Overview
86
+
87
+ ONNX Script provides a few major capabilities for authoring and debugging
88
+ ONNX models and functions:
89
+
90
+ * A converter which translates a Python ONNX Script function into an
91
+ ONNX graph, accomplished by traversing the [Python Abstract Syntax Tree][python-ast] to build an ONNX graph equivalent of the function.
92
+
93
+ * A converter that operates inversely, translating ONNX models and
94
+ functions into ONNX Script. This capability can be used to fully round-trip
95
+ ONNX Script ↔ ONNX graph.
96
+
97
+ * A runtime shim that allows such functions to be evaluated
98
+ (in an "eager mode"). This functionality currently relies on
99
+ [ONNX Runtime][onnx-runtime] for executing every [ONNX Operator][onnx-ops],
100
+ and there is a Python-only reference runtime for ONNX underway that
101
+ will also be supported.
102
+
103
+ Note that the runtime is intended to help understand and debug function definitions. Performance is not a goal here.
104
+
105
+ ## Installing ONNX Script
106
+
107
+ ```bash
108
+ pip install --upgrade onnxscript
109
+ ```
110
+
111
+ ### Install for Development
112
+
113
+ ```bash
114
+ git clone https://github.com/microsoft/onnxscript
115
+ cd onnxscript
116
+ pip install -r requirements-dev.txt
117
+ pip install -e .
118
+ ```
119
+
120
+ ### Run Unit Tests
121
+
122
+ ```bash
123
+ pytest .
124
+ ```
125
+
126
+ ## Example
127
+
128
+ ```python update-readme
129
+ import onnx
130
+
131
+ # We use ONNX opset 15 to define the function below.
132
+ from onnxscript import FLOAT, script
133
+ from onnxscript import opset15 as op
134
+
135
+
136
+ # We use the script decorator to indicate that
137
+ # this is meant to be translated to ONNX.
138
+ @script()
139
+ def onnx_hardmax(X, axis: int):
140
+ """Hardmax is similar to ArgMax, with the result being encoded OneHot style."""
141
+
142
+ # The type annotation on X indicates that it is a float tensor of
143
+ # unknown rank. The type annotation on axis indicates that it will
144
+ # be treated as an int attribute in ONNX.
145
+ #
146
+ # Invoke ONNX opset 15 op ArgMax.
147
+ # Use unnamed arguments for ONNX input parameters, and named
148
+ # arguments for ONNX attribute parameters.
149
+ argmax = op.ArgMax(X, axis=axis, keepdims=False)
150
+ xshape = op.Shape(X, start=axis)
151
+ # use the Constant operator to create constant tensors
152
+ zero = op.Constant(value_ints=[0])
153
+ depth = op.GatherElements(xshape, zero)
154
+ empty_shape = op.Constant(value_ints=[0])
155
+ depth = op.Reshape(depth, empty_shape)
156
+ values = op.Constant(value_ints=[0, 1])
157
+ cast_values = op.CastLike(values, X)
158
+ return op.OneHot(argmax, depth, cast_values, axis=axis)
159
+
160
+
161
+ # We use the script decorator to indicate that
162
+ # this is meant to be translated to ONNX.
163
+ @script()
164
+ def sample_model(X: FLOAT[64, 128], Wt: FLOAT[128, 10], Bias: FLOAT[10]) -> FLOAT[64, 10]:
165
+ matmul = op.MatMul(X, Wt) + Bias
166
+ return onnx_hardmax(matmul, axis=1)
167
+
168
+
169
+ # onnx_model is an in-memory ModelProto
170
+ onnx_model = sample_model.to_model_proto()
171
+
172
+ # Save the ONNX model at a given path
173
+ onnx.save(onnx_model, "sample_model.onnx")
174
+
175
+ # Check the model
176
+ try:
177
+ onnx.checker.check_model(onnx_model)
178
+ except onnx.checker.ValidationError as e:
179
+ print(f"The model is invalid: {e}")
180
+ else:
181
+ print("The model is valid!")
182
+ ```
183
+
184
+ The decorator parses the code of the function, converting it into an
185
+ intermediate representation. If it fails, it produces an error message
186
+ indicating the line where the error was detected. If it succeeds, the
187
+ intermediate representation can be converted into an ONNX graph
188
+ structure of type `FunctionProto`:
189
+
190
+ * `Hardmax.to_function_proto()` returns a `FunctionProto`
191
+
192
+ ### Eager Mode Evaluation
193
+
194
+ Eager mode is mostly used to debug and validate that intermediate results
195
+ are as expected. The function defined above can be called as below,
196
+ executing in an eager-evaluation mode:
197
+
198
+ ```python
199
+ import numpy as np
200
+
201
+ v = np.array([[0, 1], [2, 3]], dtype=np.float32)
202
+ result = Hardmax(v)
203
+ ```
204
+
205
+ More examples can be found in the [docs/examples](docs/examples) directory.
206
+
207
+ ## ONNX IR
208
+
209
+ An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
210
+
211
+ ### Features
212
+
213
+ * **Full ONNX spec support:** all valid models representable by ONNX protobuf,
214
+ and a subset of invalid models (so you can load and fix them).
215
+ * **Low memory footprint:** mmap'ed external tensors; unified interface for
216
+ ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size
217
+ limitation. Zero copies.
218
+ * **Straightforward access patterns:** Access value information and traverse the
219
+ graph topology at ease.
220
+ * **Robust mutation:** Create as many iterators as you like on the graph while mutating it.
221
+ * **Speed:** Performant graph manipulation, serialization/deserialization to Protobuf.
222
+ * **Pythonic and familiar APIs:** Classes define Pythonic apis and still map to
223
+ ONNX protobuf concepts in an intuitive way.
224
+
225
+ ## ONNX Script Tools
226
+
227
+ ### ONNX Optimizer
228
+
229
+ The ONNX Script Optimizer tool provides the user with the functionality to optimize an ONNX model by performing optimizations and clean-ups such as constant folding, dead code elimination, etc. In order to utilize the optimizer tool:
230
+
231
+ ```python
232
+ import onnxscript
233
+
234
+ onnxscript.optimizer.optimize(onnx_model)
235
+ ```
236
+
237
+ For a detailed summary of all the optimizations applied by the optimizer call, refer to the tutorial [Optimizing a Model using the Optimizer](https://onnxscript.ai/tutorial/optimizer/optimize.html)
238
+
239
+ ### ONNX Rewriter
240
+
241
+ The ONNX Rewriter tool provides the user with the functionality to replace certain patterns in an ONNX graph with another pattern based on user-defined rewrite rules. The rewriter tools allows two different methods in which patterns in the graph can be rewritten.
242
+
243
+ ### Pattern-based rewriting
244
+
245
+ For this style of rewriting, the user provides a `target_pattern` that is to be replaced, a `replacement_pattern` and a `match_condition` (pattern rewrite will occur only if the match condition is satisfied). A simple example on how to use the pattern-based rewriting tool is as follows:
246
+
247
+ ```python
248
+ from onnxscript.rewriter import pattern
249
+
250
+ # The target pattern
251
+ def erf_gelu_pattern(op, x):
252
+ return 0.5 * (x * (op.Erf(x / math.sqrt(2)) + 1.0))
253
+
254
+ def erf_gelu_pattern_2(op, x):
255
+ return (x * (op.Erf(x / math.sqrt(2)) + 1.0)) * 0.5
256
+
257
+ # The replacement pattern
258
+ def gelu(op, x: ir.Value):
259
+ return op.Gelu(x, domain="com.microsoft")
260
+
261
+ # Create multiple rules
262
+ rule1 = pattern.RewriteRule(
263
+ erf_gelu_pattern, # Target Pattern
264
+ gelu, # Replacement
265
+ )
266
+ rule2 = pattern.RewriteRule(
267
+ erf_gelu_pattern_2, # Target Pattern
268
+ gelu, # Replacement
269
+ )
270
+ # Create a Rewrite Rule Set with multiple rules.
271
+ rewrite_rule_set = pattern.RewriteRuleSet([rule1, rule2])
272
+ # Apply rewrites
273
+ model_with_rewrite_applied = onnxscript.rewriter.rewrite(
274
+ model, # Original ONNX Model
275
+ pattern_rewrite_rules=rewrite_rule_set,
276
+ )
277
+ return model_with_rewrite_applied
278
+ ```
279
+
280
+ For a detailed tutorial on how to create target_pattern, replacement_pattern and match_condition blocks in order to utilize the pattern-based rewriter, refer to the tutorial [Pattern-based Rewrite Using Rules](https://onnxscript.ai/tutorial/rewriter/rewrite_patterns.html)
281
+
282
+ ### Function-based rewriting
283
+
284
+ This style of rewriting matches a `FUNCTION_KEYWORD` and `PACKAGE_NAME` provided by the user to an existing function within the graph and replaces it with a new function provided by the user.
285
+
286
+ ## Development Guidelines
287
+
288
+ Every change impacting the converter or the eager evaluation must be
289
+ unit tested with class `OnnxScriptTestCase` to ensure both systems do
290
+ return the same results with the same inputs.
291
+
292
+ ### Coding Style
293
+
294
+ We use `ruff`, `black`, `isort`, and `mypy` etc. to check code formatting and use `lintrunner` to run all linters.
295
+ You can install the dependencies and initialize with
296
+
297
+ ```sh
298
+ pip install lintrunner lintrunner-adapters
299
+ lintrunner init
300
+ ```
301
+
302
+ This will install lintrunner on your system and download all the necessary dependencies to run linters locally.
303
+ If you want to see what lintrunner init will install, run `lintrunner init --dry-run`.
304
+
305
+ To lint local changes:
306
+
307
+ ```bash
308
+ lintrunner
309
+ ```
310
+
311
+ To format files:
312
+
313
+ ```bash
314
+ lintrunner f
315
+ ```
316
+
317
+ To lint all files:
318
+
319
+ ```bash
320
+ lintrunner --all-files
321
+ ```
322
+
323
+ Use `--output oneline` to produce a compact list of lint errors, useful when
324
+ there are many errors to fix.
325
+
326
+ See all available options with `lintrunner -h`.
327
+
328
+ To read more about lintrunner, see [wiki](https://github.com/pytorch/pytorch/wiki/lintrunner).
329
+ To update an existing linting rule or create a new one, modify `.lintrunner.toml` or create a
330
+ new adapter following examples in https://github.com/justinchuby/lintrunner-adapters.
331
+
332
+ ## Contributing
333
+
334
+ We're always looking for your help to improve the product (bug fixes, new features, documentation, etc). Currently ONNX Script is under early and heavy development, so we encourage proposing any major changes by [filing an issue](https://github.com/microsoft/onnxscript/issues) to discuss your idea with the team first.
335
+
336
+ ### Report a Security Issue
337
+
338
+ **Please do not report security vulnerabilities through public GitHub issues.**
339
+
340
+ Please refer to our guidance on filing [Security Issues](SECURITY.md).
341
+
342
+ ### Licensing Guidelines
343
+
344
+ This project welcomes contributions and suggestions. Most contributions require you to
345
+ agree to a Contributor License Agreement (CLA) declaring that you have the right to,
346
+ and actually do, grant us the rights to use your contribution. For details, visit
347
+ https://cla.microsoft.com.
348
+
349
+ When you submit a pull request, a CLA-bot will automatically determine whether you need
350
+ to provide a CLA and decorate the PR appropriately (e.g., label, comment). Simply follow the
351
+ instructions provided by the bot. You will only need to do this once across all repositories using our CLA.
352
+
353
+ ### Code of Conduct
354
+
355
+ This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
356
+ For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
357
+ or contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.
358
+
359
+ ## Trademarks
360
+
361
+ This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft
362
+ trademarks or logos is subject to and must follow
363
+ [Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general).
364
+ Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship.
365
+ Any use of third-party trademarks or logos is subject to those third-party's policies.
366
+
367
+ [python-ast]: https://docs.python.org/3/library/ast.html
368
+ [onnx-runtime]: https://onnxruntime.ai
369
+ [onnx-ops]: https://github.com/onnx/onnx/blob/main/docs/Operators.md
370
+ [onnxfns1A.py]: https://github.com/microsoft/onnxscript/blob/main/onnxscript/tests/models/onnxfns1A.py