Skip to content

Commit

Permalink
Add Q/DQ operators to ORT's symbolic shape inference script
Browse files Browse the repository at this point in the history
  • Loading branch information
adrianlizarraga committed Jan 30, 2024
1 parent ffc3431 commit f7e14db
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions onnxruntime/python/tools/symbolic_shape_infer.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ def __init__(self, int_max, auto_merge, guess_output_rank, verbose, prefix=""):
"ConstantOfShape": self._infer_ConstantOfShape,
"Conv": self._infer_Conv,
"CumSum": self._pass_on_shape_and_type,
"DequantizeLinear": self._infer_DequantizeLinear,
"Div": self._infer_symbolic_compute_ops,
"Einsum": self._infer_Einsum,
"Expand": self._infer_Expand,
Expand All @@ -163,6 +164,7 @@ def __init__(self, int_max, auto_merge, guess_output_rank, verbose, prefix=""):
"NonZero": self._infer_NonZero,
"OneHot": self._infer_OneHot,
"Pad": self._infer_Pad,
"QuantizeLinear": self._infer_QuantizeLinear,
"Range": self._infer_Range,
"Reciprocal": self._pass_on_shape_and_type,
"ReduceSum": self._infer_ReduceSum,
Expand Down Expand Up @@ -979,6 +981,12 @@ def _infer_NhwcConv(self, node): # noqa: N802
)
)

def _infer_DequantizeLinear(self, node): # noqa: N802
self._propagate_shape_and_type(node)

def _infer_QuantizeLinear(self, node): # noqa: N802
self._propagate_shape_and_type(node)

def _infer_Einsum(self, node): # noqa: N802
# ref:https://github.com/onnx/onnx/blob/623dfaa0151b2e4ce49779c3ec31cbd78c592b80/onnx/defs/math/defs.cc#L3275
equation = get_attribute(node, "equation")
Expand Down

0 comments on commit f7e14db

Please sign in to comment.