onnxscript 0.1.0__py3-none-any.whl

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 (176) hide show
  1. onnxscript/__init__.py +131 -0
  2. onnxscript/_framework_apis/__init__.py +3 -0
  3. onnxscript/_framework_apis/torch_2_5.py +117 -0
  4. onnxscript/_framework_apis/torch_2_6.py +45 -0
  5. onnxscript/_internal/__init__.py +0 -0
  6. onnxscript/_internal/analysis.py +229 -0
  7. onnxscript/_internal/ast_utils.py +64 -0
  8. onnxscript/_internal/autocast.py +250 -0
  9. onnxscript/_internal/deprecation.py +78 -0
  10. onnxscript/_internal/param_manipulation.py +148 -0
  11. onnxscript/_internal/runtime_typing.py +43 -0
  12. onnxscript/_internal/utils.py +99 -0
  13. onnxscript/_internal/version_utils.py +118 -0
  14. onnxscript/_legacy_ir/__init__.py +341 -0
  15. onnxscript/_legacy_ir/visitor.py +937 -0
  16. onnxscript/_thirdparty/asciichartpy.py +313 -0
  17. onnxscript/backend/__init__.py +2 -0
  18. onnxscript/backend/onnx_backend.py +303 -0
  19. onnxscript/backend/onnx_export.py +885 -0
  20. onnxscript/converter.py +1470 -0
  21. onnxscript/evaluator.py +619 -0
  22. onnxscript/function_libs/tools/torch_lib/deduce_type_constraints.py +403 -0
  23. onnxscript/function_libs/tools/torch_lib/generate_aten_signatures.py +333 -0
  24. onnxscript/function_libs/tools/torch_lib/generate_prims_signatures.py +331 -0
  25. onnxscript/function_libs/torch_lib/__init__.py +12 -0
  26. onnxscript/function_libs/torch_lib/_constants.py +5 -0
  27. onnxscript/function_libs/torch_lib/_flags.py +58 -0
  28. onnxscript/function_libs/torch_lib/graph_building/__init__.py +56 -0
  29. onnxscript/function_libs/torch_lib/graph_building/_graph_building_ir.py +723 -0
  30. onnxscript/function_libs/torch_lib/graph_building/_graph_building_torch.py +1125 -0
  31. onnxscript/function_libs/torch_lib/ops/__init__.py +27 -0
  32. onnxscript/function_libs/torch_lib/ops/common.py +80 -0
  33. onnxscript/function_libs/torch_lib/ops/core.py +8935 -0
  34. onnxscript/function_libs/torch_lib/ops/fft.py +385 -0
  35. onnxscript/function_libs/torch_lib/ops/linalg.py +399 -0
  36. onnxscript/function_libs/torch_lib/ops/nested.py +25 -0
  37. onnxscript/function_libs/torch_lib/ops/nn.py +2713 -0
  38. onnxscript/function_libs/torch_lib/ops/prims.py +850 -0
  39. onnxscript/function_libs/torch_lib/ops/quantized_decomposed.py +63 -0
  40. onnxscript/function_libs/torch_lib/ops/sparse.py +23 -0
  41. onnxscript/function_libs/torch_lib/ops/special.py +387 -0
  42. onnxscript/function_libs/torch_lib/ops/vision.py +25 -0
  43. onnxscript/function_libs/torch_lib/registration.py +151 -0
  44. onnxscript/function_libs/torch_lib/tensor_typing.py +74 -0
  45. onnxscript/ir/__init__.py +153 -0
  46. onnxscript/ir/_convenience.py +447 -0
  47. onnxscript/ir/_core.py +3119 -0
  48. onnxscript/ir/_display.py +49 -0
  49. onnxscript/ir/_enums.py +163 -0
  50. onnxscript/ir/_graph_comparison.py +23 -0
  51. onnxscript/ir/_io.py +97 -0
  52. onnxscript/ir/_linked_list.py +276 -0
  53. onnxscript/ir/_metadata.py +44 -0
  54. onnxscript/ir/_name_authority.py +72 -0
  55. onnxscript/ir/_polyfill.py +25 -0
  56. onnxscript/ir/_protocols.py +598 -0
  57. onnxscript/ir/_schemas.py +548 -0
  58. onnxscript/ir/_tape.py +120 -0
  59. onnxscript/ir/_type_casting.py +106 -0
  60. onnxscript/ir/convenience.py +32 -0
  61. onnxscript/ir/external_data.py +396 -0
  62. onnxscript/ir/passes/__init__.py +33 -0
  63. onnxscript/ir/passes/_pass_infra.py +172 -0
  64. onnxscript/ir/serde.py +1620 -0
  65. onnxscript/ir/tensor_adapters.py +122 -0
  66. onnxscript/ir/traversal.py +82 -0
  67. onnxscript/irbuilder.py +542 -0
  68. onnxscript/main.py +167 -0
  69. onnxscript/onnx_opset/__init__.py +232 -0
  70. onnxscript/onnx_opset/_impl/opset1.py +4100 -0
  71. onnxscript/onnx_opset/_impl/opset10.py +1227 -0
  72. onnxscript/onnx_opset/_impl/opset11.py +4013 -0
  73. onnxscript/onnx_opset/_impl/opset12.py +1078 -0
  74. onnxscript/onnx_opset/_impl/opset13.py +3924 -0
  75. onnxscript/onnx_opset/_impl/opset14.py +999 -0
  76. onnxscript/onnx_opset/_impl/opset15.py +604 -0
  77. onnxscript/onnx_opset/_impl/opset16.py +1255 -0
  78. onnxscript/onnx_opset/_impl/opset17.py +561 -0
  79. onnxscript/onnx_opset/_impl/opset18.py +1803 -0
  80. onnxscript/onnx_opset/_impl/opset19.py +1942 -0
  81. onnxscript/onnx_opset/_impl/opset2.py +218 -0
  82. onnxscript/onnx_opset/_impl/opset20.py +675 -0
  83. onnxscript/onnx_opset/_impl/opset21.py +1976 -0
  84. onnxscript/onnx_opset/_impl/opset22.py +2588 -0
  85. onnxscript/onnx_opset/_impl/opset3.py +199 -0
  86. onnxscript/onnx_opset/_impl/opset4.py +77 -0
  87. onnxscript/onnx_opset/_impl/opset5.py +84 -0
  88. onnxscript/onnx_opset/_impl/opset6.py +944 -0
  89. onnxscript/onnx_opset/_impl/opset7.py +1243 -0
  90. onnxscript/onnx_opset/_impl/opset8.py +444 -0
  91. onnxscript/onnx_opset/_impl/opset9.py +1485 -0
  92. onnxscript/onnx_opset/_impl/opset_ai_onnx_ml1.py +974 -0
  93. onnxscript/onnx_opset/_impl/opset_ai_onnx_ml2.py +112 -0
  94. onnxscript/onnx_opset/_impl/opset_ai_onnx_ml3.py +308 -0
  95. onnxscript/onnx_opset/_impl/opset_ai_onnx_ml4.py +129 -0
  96. onnxscript/onnx_opset/_impl/opset_ai_onnx_ml5.py +158 -0
  97. onnxscript/onnx_opset/_impl/opset_ai_onnx_preview_training1.py +577 -0
  98. onnxscript/onnx_types.py +229 -0
  99. onnxscript/optimizer/__init__.py +39 -0
  100. onnxscript/optimizer/_constant_folding.py +1083 -0
  101. onnxscript/optimizer/_inliner.py +312 -0
  102. onnxscript/optimizer/_legacy/_optimizer.py +98 -0
  103. onnxscript/optimizer/_legacy/_remove_unused_proto.py +144 -0
  104. onnxscript/optimizer/_legacy/_simple_function_folding.py +243 -0
  105. onnxscript/optimizer/_legacy/constant_folding.py +293 -0
  106. onnxscript/optimizer/_legacy/evaluator.py +439 -0
  107. onnxscript/optimizer/_optimizer.py +61 -0
  108. onnxscript/optimizer/_remove_unused.py +106 -0
  109. onnxscript/optimizer/_remove_unused_function.py +72 -0
  110. onnxscript/py.typed +1 -0
  111. onnxscript/rewriter/__init__.py +56 -0
  112. onnxscript/rewriter/_ir_utils.py +111 -0
  113. onnxscript/rewriter/broadcast_to_matmul.py +180 -0
  114. onnxscript/rewriter/cast_constant_of_shape.py +50 -0
  115. onnxscript/rewriter/collapse_slices.py +141 -0
  116. onnxscript/rewriter/erfgelu.py +27 -0
  117. onnxscript/rewriter/function_rule.py +232 -0
  118. onnxscript/rewriter/gemm_to_matmul_add.py +21 -0
  119. onnxscript/rewriter/generic_pattern.py +700 -0
  120. onnxscript/rewriter/llama_rule_sets.py +287 -0
  121. onnxscript/rewriter/no_op.py +56 -0
  122. onnxscript/rewriter/onnxruntime/__init__.py +50 -0
  123. onnxscript/rewriter/onnxruntime/bfloat16_utils/bfloat16_converter.py +99 -0
  124. onnxscript/rewriter/onnxruntime/fused_matmul_rule_sets.py +177 -0
  125. onnxscript/rewriter/onnxruntime/group_normalization_merge_silu.py +64 -0
  126. onnxscript/rewriter/onnxruntime/instance_to_group_normalization.py +155 -0
  127. onnxscript/rewriter/onnxruntime/softmax.py +63 -0
  128. onnxscript/rewriter/onnxruntime/transformers/__init__.py +21 -0
  129. onnxscript/rewriter/onnxruntime/transformers/biassplitgelu.py +31 -0
  130. onnxscript/rewriter/onnxruntime/transformers/fastgelu.py +29 -0
  131. onnxscript/rewriter/onnxruntime/transformers/layernorm.py +47 -0
  132. onnxscript/rewriter/onnxruntime/transformers/multihead_attention.py +715 -0
  133. onnxscript/rewriter/ort_fusions/__init__.py +9 -0
  134. onnxscript/rewriter/ort_fusions/_core.py +28 -0
  135. onnxscript/rewriter/ort_fusions/_smollm_1.py +253 -0
  136. onnxscript/rewriter/ort_fusions/_smollm_2.py +467 -0
  137. onnxscript/rewriter/ort_fusions/_test_models.py +122 -0
  138. onnxscript/rewriter/ort_fusions/_test_utils.py +42 -0
  139. onnxscript/rewriter/ort_fusions/cos_sin_cache.py +154 -0
  140. onnxscript/rewriter/ort_fusions/gqa.py +156 -0
  141. onnxscript/rewriter/ort_fusions/mha.py +198 -0
  142. onnxscript/rewriter/ort_fusions/rms_normalization.py +95 -0
  143. onnxscript/rewriter/ort_fusions/rotary_embedding.py +64 -0
  144. onnxscript/rewriter/ort_fusions/sdpa.py +75 -0
  145. onnxscript/rewriter/ort_fusions/skip_normalization.py +46 -0
  146. onnxscript/rewriter/pattern.py +1714 -0
  147. onnxscript/rewriter/testing.py +77 -0
  148. onnxscript/sourceinfo.py +59 -0
  149. onnxscript/tensor.py +227 -0
  150. onnxscript/testing/__init__.py +482 -0
  151. onnxscript/tools/__init__.py +4 -0
  152. onnxscript/tools/benchmark/__init__.py +23 -0
  153. onnxscript/tools/benchmark/benchmark_helpers.py +783 -0
  154. onnxscript/tools/benchmark/benchmark_run.py +140 -0
  155. onnxscript/tools/benchmark/export_model.py +207 -0
  156. onnxscript/tools/benchmark/export_model_batch.py +146 -0
  157. onnxscript/tools/memory_peak.py +244 -0
  158. onnxscript/tools/training_helper.py +50 -0
  159. onnxscript/tools/transformers_models/__init__.py +190 -0
  160. onnxscript/tools/transformers_models/llama.py +168 -0
  161. onnxscript/tools/transformers_models/mistral.py +238 -0
  162. onnxscript/tools/transformers_models/phi.py +248 -0
  163. onnxscript/tools/transformers_models/phi3.py +259 -0
  164. onnxscript/type_annotation.py +281 -0
  165. onnxscript/utils/__init__.py +0 -0
  166. onnxscript/utils/evaluation_utils.py +56 -0
  167. onnxscript/utils/timing_utils.py +33 -0
  168. onnxscript/utils/utils.py +84 -0
  169. onnxscript/values.py +790 -0
  170. onnxscript/version_converter/__init__.py +21 -0
  171. onnxscript/version_converter/_version_converter.py +314 -0
  172. onnxscript-0.1.0.dist-info/LICENSE +21 -0
  173. onnxscript-0.1.0.dist-info/METADATA +370 -0
  174. onnxscript-0.1.0.dist-info/RECORD +176 -0
  175. onnxscript-0.1.0.dist-info/WHEEL +5 -0
  176. onnxscript-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,675 @@
1
+ # --------------------------------------------------------------------------
2
+ # ⚠️ WARNING - AUTO-GENERATED CODE - DO NOT EDIT ⚠️
3
+ # ⚙️ Generated by 'python -m opgen'
4
+ # --------------------------------------------------------------------------
5
+ # Copyright (c) Microsoft Corporation. All rights reserved.
6
+ # Licensed under the MIT License.
7
+ # --------------------------------------------------------------------------
8
+ # pylint: disable=W0221,W0222,R0901,W0237
9
+ # mypy: disable-error-code=override
10
+ # ruff: noqa: N801,E741
11
+ # ruff: noqa: D214,D402,D405,D411,D412,D416,D417
12
+ # --------------------------------------------------------------------------
13
+
14
+ from __future__ import annotations
15
+
16
+ from typing import Optional, Tuple, TypeVar, Union
17
+
18
+ from onnx import TensorProto
19
+ from onnx.defs import get_schema
20
+ from typing_extensions import TypeAlias
21
+
22
+ from onnxscript.onnx_opset._impl.opset19 import Opset19
23
+ from onnxscript.onnx_types import (
24
+ BFLOAT16,
25
+ BOOL,
26
+ COMPLEX64,
27
+ COMPLEX128,
28
+ DOUBLE,
29
+ FLOAT,
30
+ FLOAT8E4M3FN,
31
+ FLOAT8E4M3FNUZ,
32
+ FLOAT8E5M2,
33
+ FLOAT8E5M2FNUZ,
34
+ FLOAT16,
35
+ INT8,
36
+ INT16,
37
+ INT32,
38
+ INT64,
39
+ STRING,
40
+ UINT8,
41
+ UINT16,
42
+ UINT32,
43
+ UINT64,
44
+ )
45
+ from onnxscript.values import Op, Opset
46
+
47
+
48
+ class Opset20(Opset19):
49
+ def __new__(cls):
50
+ return Opset.__new__(cls, "", 20)
51
+
52
+ T1_AffineGrid = TypeVar("T1_AffineGrid", BFLOAT16, DOUBLE, FLOAT, FLOAT16)
53
+
54
+ T2_AffineGrid: TypeAlias = INT64
55
+
56
+ def AffineGrid(
57
+ self, theta: T1_AffineGrid, size: T2_AffineGrid, *, align_corners: int = 0
58
+ ) -> T1_AffineGrid:
59
+ r"""[🌐 AffineGrid(20)](https://onnx.ai/onnx/operators/onnx__AffineGrid.html#affinegrid-20 "Online Documentation")
60
+
61
+
62
+ Generates a 2D or 3D flow field (sampling grid), given a batch of affine matrices theta
63
+ (https://pytorch.org/docs/stable/generated/torch.nn.functional.affine_grid.html).
64
+ An affine matrix `theta` is applied to a position tensor represented in its homogeneous expression. Here is an example in 3D:
65
+ ::
66
+
67
+ [r00, r01, r02, t0] [x] [x']
68
+ [r10, r11, r12, t1] * [y] = [y']
69
+ [r20, r21, r22, t2] [z] [z']
70
+ [0, 0, 0, 1 ] [1] [1 ]
71
+
72
+
73
+ where `(x, y, z)` is the position in the original space, `(x', y', z')` is the position in the output space.
74
+ The last row is always `[0, 0, 0, 1]` and is not stored in the affine matrix. Therefore we have `theta` of shape `(N, 2, 3)` for 2D or `(N, 3, 4)` for 3D.
75
+
76
+ Input `size` is used to define grid of positions evenly spaced in the original 2D or 3D space, with dimensions ranging from `-1` to `1`.
77
+ The output `grid` contains positions in the output space.
78
+
79
+ When `align_corners=1`, consider `-1` and `1` to refer to the centers of the corner pixels (mark `v` in illustration).
80
+ ::
81
+
82
+ v v v v
83
+ |-------------------|------------------|
84
+ -1 0 1
85
+
86
+
87
+ When `align_corners=0`, consider `-1` and `1` to refer to the outer edge of the corner pixels.
88
+ ::
89
+
90
+ v v v v
91
+ |------------------|-------------------|
92
+ -1 0 1
93
+
94
+
95
+
96
+
97
+ Args:
98
+ theta: (non-differentiable) input batch of affine matrices with shape (N, 2,
99
+ 3) for 2D or (N, 3, 4) for 3D
100
+
101
+ size: (non-differentiable) the target output image size (N, C, H, W) for 2D
102
+ or (N, C, D, H, W) for 3D
103
+
104
+ align_corners: if align_corners=1, consider -1 and 1 to refer to the centers
105
+ of the corner pixels. if align_corners=0, consider -1 and 1 to refer to
106
+ the outer edge the corner pixels.
107
+ """
108
+
109
+ schema = get_schema("AffineGrid", 20, "")
110
+ op = Op(self, "AffineGrid", schema)
111
+ return op(*self._prepare_inputs(schema, theta, size), align_corners=align_corners)
112
+
113
+ T1_ConstantOfShape: TypeAlias = INT64
114
+
115
+ T2_ConstantOfShape: TypeAlias = Union[
116
+ BFLOAT16,
117
+ BOOL,
118
+ DOUBLE,
119
+ FLOAT,
120
+ FLOAT16,
121
+ FLOAT8E4M3FN,
122
+ FLOAT8E4M3FNUZ,
123
+ FLOAT8E5M2,
124
+ FLOAT8E5M2FNUZ,
125
+ INT16,
126
+ INT32,
127
+ INT64,
128
+ INT8,
129
+ UINT16,
130
+ UINT32,
131
+ UINT64,
132
+ UINT8,
133
+ ]
134
+
135
+ def ConstantOfShape(
136
+ self, input: T1_ConstantOfShape, *, value: Optional[TensorProto] = None
137
+ ) -> T2_ConstantOfShape:
138
+ r"""[🌐 ConstantOfShape(20)](https://onnx.ai/onnx/operators/onnx__ConstantOfShape.html#constantofshape-20 "Online Documentation")
139
+
140
+
141
+ Generate a tensor with given value and shape.
142
+
143
+
144
+ Args:
145
+ input: 1D tensor. The shape of the expected output tensor. If empty tensor
146
+ is given, the output would be a scalar. All values must be >= 0.
147
+
148
+ value: (Optional) The value of the output elements.Should be a one-element
149
+ tensor. If not specified, it defaults to a tensor of value 0 and
150
+ datatype float32
151
+ """
152
+
153
+ schema = get_schema("ConstantOfShape", 20, "")
154
+ op = Op(self, "ConstantOfShape", schema)
155
+ return op(*self._prepare_inputs(schema, input), value=value)
156
+
157
+ T1_DFT = TypeVar("T1_DFT", BFLOAT16, DOUBLE, FLOAT, FLOAT16)
158
+
159
+ T2_DFT = TypeVar("T2_DFT", INT32, INT64)
160
+
161
+ def DFT(
162
+ self,
163
+ input: T1_DFT,
164
+ dft_length: Optional[T2_DFT] = None,
165
+ axis: Optional[INT64] = None,
166
+ *,
167
+ inverse: int = 0,
168
+ onesided: int = 0,
169
+ ) -> T1_DFT:
170
+ r"""[🌐 DFT(20)](https://onnx.ai/onnx/operators/onnx__DFT.html#dft-20 "Online Documentation")
171
+
172
+ Computes the discrete Fourier Transform (DFT) of the input.
173
+
174
+ Assuming the input has shape `[M, N]`, where `N` is the dimension over which the
175
+ DFT is computed and `M` denotes the conceptual "all other dimensions,"
176
+ the DFT `y[m, k]` of shape `[M, N]` is defined as
177
+
178
+ $$y[m, k] = \sum_{n=0}^{N-1} e^{-2 \pi j \frac{k n}{N} } x[m, n] ,$$
179
+
180
+ and the inverse transform is defined as
181
+
182
+ $$x[m, n] = \frac{1}{N} \sum_{k=0}^{N-1} e^{2 \pi j \frac{k n}{N} } y[m, k] ,$$
183
+
184
+ where $j$ is the imaginary unit.
185
+
186
+ The actual shape of the output is specified in the "output" section.
187
+
188
+ Reference: https://docs.scipy.org/doc/scipy/tutorial/fft.html
189
+
190
+
191
+ Args:
192
+ input: (non-differentiable) For real input, the following shape is expected:
193
+ `[signal_dim0][signal_dim1][signal_dim2]...[signal_dimN][1]`. For
194
+ complex input, the following shape is expected:
195
+ `[signal_dim0][signal_dim1][signal_dim2]...[signal_dimN][2]`. The final
196
+ dimension represents the real and imaginary parts of the value in that
197
+ order.
198
+
199
+ dft_length: (optional, non-differentiable) The length of the signal as a
200
+ scalar. If greater than the axis dimension, the signal will be
201
+ zero-padded up to `dft_length`. If less than the axis dimension, only
202
+ the first `dft_length` values will be used as the signal.
203
+
204
+ axis: (optional, non-differentiable) The axis as a scalar on which to
205
+ perform the DFT. Default is `-2` (last signal axis). Negative value
206
+ means counting dimensions from the back. Accepted range is $[-r, -2]
207
+ \cup [0, r-2]$ where `r = rank(input)`. The last dimension is for
208
+ representing complex numbers and thus is an invalid axis.
209
+
210
+ inverse: Whether to perform the inverse discrete Fourier Transform. Default
211
+ is 0, which corresponds to `false`.
212
+
213
+ onesided: If `onesided` is `1` and input is real, only values for `k` in
214
+ `[0, 1, 2, ..., floor(n_fft/2) + 1]` are returned because the
215
+ real-to-complex Fourier transform satisfies the conjugate symmetry,
216
+ i.e., `X[m, k] = X[m, n_fft-k]*`, where `m` denotes "all other
217
+ dimensions" DFT was not applied on. If the input tensor is complex,
218
+ onesided output is not possible. Value can be `0` or `1`. Default is
219
+ `0`.
220
+ """
221
+
222
+ schema = get_schema("DFT", 20, "")
223
+ op = Op(self, "DFT", schema)
224
+ return op(
225
+ *self._prepare_inputs(schema, input, dft_length, axis),
226
+ inverse=inverse,
227
+ onesided=onesided,
228
+ )
229
+
230
+ T_Gelu = TypeVar("T_Gelu", BFLOAT16, DOUBLE, FLOAT, FLOAT16)
231
+
232
+ def Gelu(self, X: T_Gelu, *, approximate: str = "none") -> T_Gelu:
233
+ r"""[🌐 Gelu(20)](https://onnx.ai/onnx/operators/onnx__Gelu.html#gelu-20 "Online Documentation")
234
+
235
+
236
+ Gelu takes one input data (Tensor<T>) and produces one
237
+ output data (Tensor<T>) where the gaussian error linear units function,
238
+ $y = 0.5 * x * (1 + erf(x/sqrt(2)))$ is applied to the tensor elementwise.
239
+ If the attribute "approximate" is set to "tanh", the function estimation,
240
+ $y = 0.5 * x * (1 + Tanh(sqrt(2/\pi) * (x + 0.044715 * x^3)))$ is used and applied
241
+ to the tensor elementwise.
242
+
243
+
244
+
245
+ Args:
246
+ X: (differentiable) Input tensor
247
+
248
+ approximate: Gelu approximation algorithm: `"tanh"`,
249
+ `"none"`(default).`"none"`: do not use approximation.`"tanh"`: use tanh
250
+ approximation.
251
+ """
252
+
253
+ schema = get_schema("Gelu", 20, "")
254
+ op = Op(self, "Gelu", schema)
255
+ return op(*self._prepare_inputs(schema, X), approximate=approximate)
256
+
257
+ T1_GridSample = TypeVar(
258
+ "T1_GridSample",
259
+ BOOL,
260
+ COMPLEX128,
261
+ COMPLEX64,
262
+ DOUBLE,
263
+ FLOAT,
264
+ FLOAT16,
265
+ INT16,
266
+ INT32,
267
+ INT64,
268
+ INT8,
269
+ STRING,
270
+ UINT16,
271
+ UINT32,
272
+ UINT64,
273
+ UINT8,
274
+ )
275
+
276
+ T2_GridSample = TypeVar("T2_GridSample", DOUBLE, FLOAT, FLOAT16)
277
+
278
+ def GridSample(
279
+ self,
280
+ X: T1_GridSample,
281
+ grid: T2_GridSample,
282
+ *,
283
+ align_corners: int = 0,
284
+ mode: str = "linear",
285
+ padding_mode: str = "zeros",
286
+ ) -> T1_GridSample:
287
+ r"""[🌐 GridSample(20)](https://onnx.ai/onnx/operators/onnx__GridSample.html#gridsample-20 "Online Documentation")
288
+
289
+
290
+ Given an input `X` and a flow-field `grid`, computes the output `Y` using `X` values and pixel locations from the `grid`.
291
+ For spatial input `X` with shape (N, C, H, W), the `grid` will have shape (N, H_out, W_out, 2),
292
+ the output `Y` will have shape (N, C, H_out, W_out). For volumetric input `X` with shape (N, C, D, H, W),
293
+ the `grid` will have shape (N, D_out, H_out, W_out, 3), the output `Y` will have shape (N, C, D_out, H_out, W_out).
294
+ More generally, for an input `X` of rank r+2 with shape (N, C, d1, d2, ..., dr),
295
+ the `grid` will have shape (N, D1_out, D2_out, ..., Dr_out, r), the output `Y` will have shape (N, C, D1_out, D2_out, ..., Dr_out).
296
+
297
+ The tensor `X` contains values at centers of square pixels (voxels, etc) locations such as (n, c, d1_in, d2_in, ..., dr_in).
298
+ The (n, d1_out, d2_out, ..., dr_out, :) values from the tensor `grid` are the normalized positions for interpolating the values
299
+ at the (n, c, d1_out, d2_out, ..., dr_out) locations from the output tensor `Y` using a specified interpolation method (the mode)
300
+ and a padding mode (for `grid` positions falling outside the 2-dimensional image).
301
+
302
+ For example, the values in `grid[n, h_out, w_out, :]` are size-2 vectors specifying normalized positions in the 2-dimensional space of `X`.
303
+ They are used to interpolate output values of `Y[n, c, h_out, w_out]`.
304
+
305
+ The GridSample operator is often used in doing grid generator and sampler in the
306
+ [Spatial Transformer Networks](https://arxiv.org/abs/1506.02025).
307
+ See also in [torch.nn.functional.grid_sample](https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html).
308
+
309
+
310
+ Args:
311
+ X: (differentiable) Input tensor of rank r+2 that has shape (N, C, D1, D2,
312
+ ..., Dr), where N is the batch size, C is the number of channels, D1,
313
+ D2, ..., Dr are the spatial dimensions.
314
+
315
+ grid: (non-differentiable) Input offset of shape (N, D1_out, D2_out, ...,
316
+ Dr_out, r), where D1_out, D2_out, ..., Dr_out are the spatial dimensions
317
+ of the grid and output, and r is the number of spatial dimensions. Grid
318
+ specifies the sampling locations normalized by the input spatial
319
+ dimensions. Therefore, it should have most values in the range of [-1,
320
+ 1]. If the grid has values outside the range of [-1, 1], the
321
+ corresponding outputs will be handled as defined by padding_mode.
322
+ Following computer vision convention, the coordinates in the length-r
323
+ location vector are listed from the innermost tensor dimension to the
324
+ outermost, the opposite of regular tensor indexing.
325
+
326
+ align_corners: If align_corners=1, the extrema (-1 and 1) are considered as
327
+ referring to the center points of the input's corner pixels (voxels,
328
+ etc.). If align_corners=0, they are instead considered as referring to
329
+ the corner points of the input's corner pixels (voxels, etc.), making
330
+ the sampling more resolution agnostic.
331
+
332
+ mode: Three interpolation modes: linear (default), nearest and cubic. The
333
+ "linear" mode includes linear and N-linear interpolation modes depending
334
+ on the number of spatial dimensions of the input tensor (i.e. linear for
335
+ 1 spatial dimension, bilinear for 2 spatial dimensions, etc.). The
336
+ "cubic" mode also includes N-cubic interpolation modes following the
337
+ same rules. The "nearest" mode rounds to the nearest even index when the
338
+ sampling point falls halfway between two indices.
339
+
340
+ padding_mode: Support padding modes for outside grid values:
341
+ `zeros`(default), `border`, `reflection`. zeros: use 0 for out-of-bound
342
+ grid locations, border: use border values for out-of-bound grid
343
+ locations, reflection: use values at locations reflected by the border
344
+ for out-of-bound grid locations. If index 0 represents the margin pixel,
345
+ the reflected value at index -1 will be the same as the value at index
346
+ 1. For location far away from the border, it will keep being reflected
347
+ until becoming in bound. If pixel location x = -3.5 reflects by border
348
+ -1 and becomes x' = 1.5, then reflects by border 1 and becomes x'' =
349
+ 0.5.
350
+ """
351
+
352
+ schema = get_schema("GridSample", 20, "")
353
+ op = Op(self, "GridSample", schema)
354
+ return op(
355
+ *self._prepare_inputs(schema, X, grid),
356
+ align_corners=align_corners,
357
+ mode=mode,
358
+ padding_mode=padding_mode,
359
+ )
360
+
361
+ T1_ImageDecoder: TypeAlias = UINT8
362
+
363
+ T2_ImageDecoder: TypeAlias = UINT8
364
+
365
+ def ImageDecoder(
366
+ self, encoded_stream: T1_ImageDecoder, *, pixel_format: str = "RGB"
367
+ ) -> T2_ImageDecoder:
368
+ r"""[🌐 ImageDecoder(20)](https://onnx.ai/onnx/operators/onnx__ImageDecoder.html#imagedecoder-20 "Online Documentation")
369
+
370
+ Loads and decodes and image from a file. If it can't decode for any reason (e.g. corrupted encoded
371
+ stream, invalid format, it will return an empty matrix).
372
+ The following image formats are supported:
373
+ * BMP
374
+ * JPEG (note: Lossless JPEG support is optional)
375
+ * JPEG2000
376
+ * TIFF
377
+ * PNG
378
+ * WebP
379
+ * Portable image format (PBM, PGM, PPM, PXM, PNM)
380
+ Decoded images follow a channel-last layout: (Height, Width, Channels).
381
+ **JPEG chroma upsampling method:**
382
+ When upsampling the chroma components by a factor of 2, the pixels are linearly interpolated so that the
383
+ centers of the output pixels are 1/4 and 3/4 of the way between input pixel centers.
384
+ When rounding, 0.5 is rounded down and up at alternative pixels locations to prevent bias towards
385
+ larger values (ordered dither pattern).
386
+ Considering adjacent input pixels A, B, and C, B is upsampled to pixels B0 and B1 so that
387
+ ::
388
+
389
+ B0 = round_half_down((1/4) * A + (3/4) * B)
390
+ B1 = round_half_up((3/4) * B + (1/4) * C)
391
+
392
+
393
+ This method, is the default chroma upsampling method in the well-established libjpeg-turbo library,
394
+ also referred as "smooth" or "fancy" upsampling.
395
+
396
+
397
+ Args:
398
+ encoded_stream: (non-differentiable) Encoded stream
399
+
400
+ pixel_format: Pixel format. Can be one of "RGB", "BGR", or "Grayscale".
401
+ """
402
+
403
+ schema = get_schema("ImageDecoder", 20, "")
404
+ op = Op(self, "ImageDecoder", schema)
405
+ return op(*self._prepare_inputs(schema, encoded_stream), pixel_format=pixel_format)
406
+
407
+ T1_IsInf = TypeVar(
408
+ "T1_IsInf",
409
+ BFLOAT16,
410
+ DOUBLE,
411
+ FLOAT,
412
+ FLOAT16,
413
+ FLOAT8E4M3FN,
414
+ FLOAT8E4M3FNUZ,
415
+ FLOAT8E5M2,
416
+ FLOAT8E5M2FNUZ,
417
+ )
418
+
419
+ T2_IsInf: TypeAlias = BOOL
420
+
421
+ def IsInf(
422
+ self, X: T1_IsInf, *, detect_negative: int = 1, detect_positive: int = 1
423
+ ) -> T2_IsInf:
424
+ r"""[🌐 IsInf(20)](https://onnx.ai/onnx/operators/onnx__IsInf.html#isinf-20 "Online Documentation")
425
+
426
+ Map infinity to true and other values to false.
427
+
428
+ Args:
429
+ X: (non-differentiable) input
430
+
431
+ detect_negative: (Optional) Whether map negative infinity to true. Default
432
+ to 1 so that negative infinity induces true. Set this attribute to 0 if
433
+ negative infinity should be mapped to false.
434
+
435
+ detect_positive: (Optional) Whether map positive infinity to true. Default
436
+ to 1 so that positive infinity induces true. Set this attribute to 0 if
437
+ positive infinity should be mapped to false.
438
+ """
439
+
440
+ schema = get_schema("IsInf", 20, "")
441
+ op = Op(self, "IsInf", schema)
442
+ return op(
443
+ *self._prepare_inputs(schema, X),
444
+ detect_negative=detect_negative,
445
+ detect_positive=detect_positive,
446
+ )
447
+
448
+ T1_IsNaN = TypeVar(
449
+ "T1_IsNaN",
450
+ BFLOAT16,
451
+ DOUBLE,
452
+ FLOAT,
453
+ FLOAT16,
454
+ FLOAT8E4M3FN,
455
+ FLOAT8E4M3FNUZ,
456
+ FLOAT8E5M2,
457
+ FLOAT8E5M2FNUZ,
458
+ )
459
+
460
+ T2_IsNaN: TypeAlias = BOOL
461
+
462
+ def IsNaN(self, X: T1_IsNaN) -> T2_IsNaN:
463
+ r"""[🌐 IsNaN(20)](https://onnx.ai/onnx/operators/onnx__IsNaN.html#isnan-20 "Online Documentation")
464
+
465
+ Returns which elements of the input are NaN.
466
+
467
+ Args:
468
+ X: (non-differentiable) input
469
+ """
470
+
471
+ schema = get_schema("IsNaN", 20, "")
472
+ op = Op(self, "IsNaN", schema)
473
+ return op(*self._prepare_inputs(schema, X))
474
+
475
+ T_ReduceMax = TypeVar(
476
+ "T_ReduceMax",
477
+ BFLOAT16,
478
+ BOOL,
479
+ DOUBLE,
480
+ FLOAT,
481
+ FLOAT16,
482
+ INT32,
483
+ INT64,
484
+ INT8,
485
+ UINT32,
486
+ UINT64,
487
+ UINT8,
488
+ )
489
+
490
+ def ReduceMax(
491
+ self,
492
+ data: T_ReduceMax,
493
+ axes: Optional[INT64] = None,
494
+ *,
495
+ keepdims: int = 1,
496
+ noop_with_empty_axes: int = 0,
497
+ ) -> T_ReduceMax:
498
+ r"""[🌐 ReduceMax(20)](https://onnx.ai/onnx/operators/onnx__ReduceMax.html#reducemax-20 "Online Documentation")
499
+
500
+
501
+ Computes the max of the input tensor's elements along the provided axes. The resulting
502
+ tensor has the same rank as the input if `keepdims` equals 1. If `keepdims` equals 0, then
503
+ the resulting tensor has the reduced dimension pruned. Input tensors of rank zero are
504
+ valid. Reduction over an empty set of values yields minus infinity (if supported by the datatype) or the minimum value of the data type otherwise.
505
+
506
+
507
+ If the input data type is Boolean, the comparison should consider `False < True`.
508
+
509
+ The above behavior is similar to numpy, with the exception that numpy defaults `keepdims`
510
+ to `False` instead of `True`.
511
+
512
+ Args:
513
+ data: (differentiable) An input tensor.
514
+
515
+ axes: (optional, non-differentiable) Optional input list of integers, along
516
+ which to reduce. The default is to reduce over all the dimensions of the
517
+ input tensor if 'noop_with_empty_axes' is false, else act as an Identity
518
+ op when 'noop_with_empty_axes' is true. Accepted range is [-r, r-1]
519
+ where r = rank(data).
520
+
521
+ keepdims: Keep the reduced dimension or not, default 1 means keep reduced
522
+ dimension.
523
+
524
+ noop_with_empty_axes: Defines behavior if 'axes' is empty. Default behavior
525
+ with 'false' is to reduce all axes. When axes is empty and this
526
+ attribute is set to true, input tensor will not be reduced,and the
527
+ output tensor would be equivalent to input tensor.
528
+ """
529
+
530
+ schema = get_schema("ReduceMax", 20, "")
531
+ op = Op(self, "ReduceMax", schema)
532
+ return op(
533
+ *self._prepare_inputs(schema, data, axes),
534
+ keepdims=keepdims,
535
+ noop_with_empty_axes=noop_with_empty_axes,
536
+ )
537
+
538
+ T_ReduceMin = TypeVar(
539
+ "T_ReduceMin",
540
+ BFLOAT16,
541
+ BOOL,
542
+ DOUBLE,
543
+ FLOAT,
544
+ FLOAT16,
545
+ INT32,
546
+ INT64,
547
+ INT8,
548
+ UINT32,
549
+ UINT64,
550
+ UINT8,
551
+ )
552
+
553
+ def ReduceMin(
554
+ self,
555
+ data: T_ReduceMin,
556
+ axes: Optional[INT64] = None,
557
+ *,
558
+ keepdims: int = 1,
559
+ noop_with_empty_axes: int = 0,
560
+ ) -> T_ReduceMin:
561
+ r"""[🌐 ReduceMin(20)](https://onnx.ai/onnx/operators/onnx__ReduceMin.html#reducemin-20 "Online Documentation")
562
+
563
+
564
+ Computes the min of the input tensor's elements along the provided axes. The resulting
565
+ tensor has the same rank as the input if `keepdims` equals 1. If `keepdims` equals 0, then
566
+ the resulting tensor has the reduced dimension pruned. Input tensors of rank zero are
567
+ valid. Reduction over an empty set of values yields plus infinity (if supported by the datatype) or the maximum value of the data type otherwise.
568
+
569
+
570
+ If the input data type is Boolean, the comparison should consider `False < True`.
571
+
572
+ The above behavior is similar to numpy, with the exception that numpy defaults `keepdims`
573
+ to `False` instead of `True`.
574
+
575
+ Args:
576
+ data: (differentiable) An input tensor.
577
+
578
+ axes: (optional, non-differentiable) Optional input list of integers, along
579
+ which to reduce. The default is to reduce over all the dimensions of the
580
+ input tensor if 'noop_with_empty_axes' is false, else act as an Identity
581
+ op when 'noop_with_empty_axes' is true. Accepted range is [-r, r-1]
582
+ where r = rank(data).
583
+
584
+ keepdims: Keep the reduced dimension or not, default 1 means keep reduced
585
+ dimension.
586
+
587
+ noop_with_empty_axes: Defines behavior if 'axes' is empty. Default behavior
588
+ with 'false' is to reduce all axes. When axes is empty and this
589
+ attribute is set to true, input tensor will not be reduced,and the
590
+ output tensor would be equivalent to input tensor.
591
+ """
592
+
593
+ schema = get_schema("ReduceMin", 20, "")
594
+ op = Op(self, "ReduceMin", schema)
595
+ return op(
596
+ *self._prepare_inputs(schema, data, axes),
597
+ keepdims=keepdims,
598
+ noop_with_empty_axes=noop_with_empty_axes,
599
+ )
600
+
601
+ T1_RegexFullMatch: TypeAlias = STRING
602
+
603
+ T2_RegexFullMatch: TypeAlias = BOOL
604
+
605
+ def RegexFullMatch(
606
+ self, X: T1_RegexFullMatch, *, pattern: Optional[str] = None
607
+ ) -> T2_RegexFullMatch:
608
+ r"""[🌐 RegexFullMatch(20)](https://onnx.ai/onnx/operators/onnx__RegexFullMatch.html#regexfullmatch-20 "Online Documentation")
609
+
610
+ RegexFullMatch performs a full regex match on each element of the input tensor. If an element fully matches the regex pattern specified as an attribute, the corresponding element in the output is True and it is False otherwise. [RE2](https://github.com/google/re2/wiki/Syntax) regex syntax is used.
611
+
612
+ Args:
613
+ X: (non-differentiable) Tensor with strings to match on.
614
+
615
+ pattern: Regex pattern to match on. This must be valid RE2 syntax.
616
+ """
617
+
618
+ schema = get_schema("RegexFullMatch", 20, "")
619
+ op = Op(self, "RegexFullMatch", schema)
620
+ return op(*self._prepare_inputs(schema, X), pattern=pattern)
621
+
622
+ T_StringConcat: TypeAlias = STRING
623
+
624
+ def StringConcat(self, X: T_StringConcat, Y: T_StringConcat) -> T_StringConcat:
625
+ r"""[🌐 StringConcat(20)](https://onnx.ai/onnx/operators/onnx__StringConcat.html#stringconcat-20 "Online Documentation")
626
+
627
+ StringConcat concatenates string tensors elementwise (with NumPy-style broadcasting support)
628
+
629
+ Args:
630
+ X: (non-differentiable) Tensor to prepend in concatenation
631
+
632
+ Y: (non-differentiable) Tensor to append in concatenation
633
+ """
634
+
635
+ schema = get_schema("StringConcat", 20, "")
636
+ op = Op(self, "StringConcat", schema)
637
+ return op(*self._prepare_inputs(schema, X, Y))
638
+
639
+ T1_StringSplit: TypeAlias = STRING
640
+
641
+ T2_StringSplit: TypeAlias = STRING
642
+
643
+ T3_StringSplit: TypeAlias = INT64
644
+
645
+ def StringSplit(
646
+ self,
647
+ X: T1_StringSplit,
648
+ *,
649
+ delimiter: Optional[str] = None,
650
+ maxsplit: Optional[int] = None,
651
+ ) -> Tuple[T2_StringSplit, T3_StringSplit]:
652
+ r"""[🌐 StringSplit(20)](https://onnx.ai/onnx/operators/onnx__StringSplit.html#stringsplit-20 "Online Documentation")
653
+
654
+ StringSplit splits a string tensor's elements into substrings based on a delimiter attribute and a maxsplit attribute.
655
+
656
+ The first output of this operator is a tensor of strings representing the substrings from splitting each input string on the `delimiter` substring. This tensor has one additional rank compared to the input tensor in order to store the substrings for each input element (where the input tensor is not empty). Note that, in order to ensure the same number of elements are present in the final dimension, this tensor will pad empty strings as illustrated in the examples below. Consecutive delimiters are not grouped together and are deemed to delimit empty strings, except if the `delimiter` is unspecified or is the empty string (""). In the case where the `delimiter` is unspecified or the empty string, consecutive whitespace characters are regarded as a single separator and leading or trailing whitespace is removed in the output.
657
+
658
+ The second output tensor represents the number of substrings generated. `maxsplit` can be used to limit the number of splits performed - after the `maxsplit`th split if the string is not fully split, the trailing suffix of input string after the final split point is also added. For elements where fewer splits are possible than specified in `maxsplit`, it has no effect.
659
+
660
+ Args:
661
+ X: (non-differentiable) Tensor of strings to split.
662
+
663
+ delimiter: Delimiter to split on. If left unset or set to the empty string
664
+ (""), the input is split on consecutive whitespace.
665
+
666
+ maxsplit: Maximum number of splits (from left to right). If left unset (or
667
+ if the number of possible splits are less than maxsplit), it will make
668
+ as many splits as possible. Note that the maximum possible number of
669
+ substrings returned with `maxsplit` specified is `maxsplit+1` since the
670
+ remaining suffix after the `maxsplit`th split is included in the output.
671
+ """
672
+
673
+ schema = get_schema("StringSplit", 20, "")
674
+ op = Op(self, "StringSplit", schema)
675
+ return op(*self._prepare_inputs(schema, X), delimiter=delimiter, maxsplit=maxsplit)