onnxslim 0.1.82__py3-none-any.whl → 0.1.84__py3-none-any.whl

onnxslim/core/shape_inference/contrib_ops/attention/attention.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Attention operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class AttentionHandler(ShapeHandler):
+    """Handler for Attention operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Attention"
+
+    def infer_shape(self, node, ctx) -> None:
+        shape = ctx.get_shape(node, 0)
+        shape_weights = ctx.get_shape(node, 1)
+        shape_bias = ctx.try_get_shape(node, 2)
+        if shape_bias is not None:
+            assert len(shape_bias) == 1
+        tripled_hidden_size = shape_bias[0] if shape_bias is not None else shape_weights[1]
+        if shape and len(shape) == 3:
+            qkv_hidden_sizes_attr = get_attribute(node, "qkv_hidden_sizes")
+            if qkv_hidden_sizes_attr is not None:
+                assert len(qkv_hidden_sizes_attr) == 3
+                shape[2] = int(qkv_hidden_sizes_attr[2])
+            elif isinstance(tripled_hidden_size, int):
+                shape[2] = int(tripled_hidden_size / 3)
+            output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+            vi = ctx.known_vi_[node.output[0]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, shape))
+
+            if len(node.output) > 1:
+                input_shape = ctx.get_shape(node, 0)
+                past_shape = ctx.get_shape(node, 4) if len(node.input) > 4 and node.input[4] else []
+                mask_shape = ctx.get_shape(node, 3) if len(node.input) > 3 and node.input[3] else []
+
+                if past_shape and len(past_shape) == 5:
+                    if mask_shape and len(mask_shape) in {2, 3}:
+                        past_shape[3] = mask_shape[-1]
+                    elif input_shape and len(input_shape) == 3:
+                        if isinstance(input_shape[1], int) and isinstance(past_shape[3], int):
+                            past_shape[3] = input_shape[1] + past_shape[3]
+                        else:
+                            past_shape[3] = f"{past_shape[3]}+{input_shape[1]}"
+                    vi = ctx.known_vi_[node.output[1]]
+                    vi.CopyFrom(helper.make_tensor_value_info(vi.name, output_dtype, past_shape))
+                else:
+                    num_heads = get_attribute(node, "num_heads")
+                    head_size = input_shape[2] // num_heads
+                    present_shape = [2, input_shape[0], num_heads, input_shape[1], head_size]
+                    vi = ctx.known_vi_[node.output[1]]
+                    vi.CopyFrom(helper.make_tensor_value_info(vi.name, output_dtype, present_shape))
+
+
+register_shape_handler(AttentionHandler())

onnxslim/core/shape_inference/contrib_ops/attention/decoder_masked_mha.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for DecoderMaskedMultiHeadAttention operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class DecoderMaskedMultiHeadAttentionHandler(ShapeHandler):
+    """Handler for DecoderMaskedMultiHeadAttention operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "DecoderMaskedMultiHeadAttention"
+
+    def infer_shape(self, node, ctx) -> None:
+        query_shape = ctx.get_shape(node, 0)
+        if query_shape is not None:
+            output_shape = query_shape
+            output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+            assert output_dtype is not None
+            vi = ctx.known_vi_[node.output[0]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+            if len(node.output) > 2 and node.output[1] and node.output[2]:
+                past_shape = ctx.try_get_shape(node, 5)
+                if past_shape is not None:
+                    vi = ctx.known_vi_[node.output[1]]
+                    vi.CopyFrom(helper.make_tensor_value_info(vi.name, output_dtype, past_shape))
+                    vi = ctx.known_vi_[node.output[2]]
+                    vi.CopyFrom(helper.make_tensor_value_info(vi.name, output_dtype, past_shape))
+
+
+register_shape_handler(DecoderMaskedMultiHeadAttentionHandler())

onnxslim/core/shape_inference/contrib_ops/attention/gated_relative_position_bias.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for GatedRelativePositionBias operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class GatedRelativePositionBiasHandler(ShapeHandler):
+    """Handler for GatedRelativePositionBias operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "GatedRelativePositionBias"
+
+    def infer_shape(self, node, ctx) -> None:
+        num_heads = get_attribute(node, "num_heads")
+        token_offset_shape = ctx.try_get_shape(node, 6)
+        if token_offset_shape is not None:
+            output_shape = [token_offset_shape[0], num_heads, token_offset_shape[1], token_offset_shape[1]]
+        else:
+            query_layer_shape = ctx.get_shape(node, 0)
+            assert query_layer_shape is not None and len(query_layer_shape) == 3
+            output_shape = [query_layer_shape[0], num_heads, query_layer_shape[1], query_layer_shape[1]]
+
+        output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+
+register_shape_handler(GatedRelativePositionBiasHandler())

onnxslim/core/shape_inference/contrib_ops/attention/longformer_attention.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for LongformerAttention operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class LongformerAttentionHandler(ShapeHandler):
+    """Handler for LongformerAttention operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "LongformerAttention"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+
+register_shape_handler(LongformerAttentionHandler())

onnxslim/core/shape_inference/contrib_ops/attention/multi_head_attention.py

@@ -0,0 +1,82 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for MultiHeadAttention operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class MultiHeadAttentionHandler(ShapeHandler):
+    """Handler for MultiHeadAttention operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "MultiHeadAttention"
+
+    def infer_shape(self, node, ctx) -> None:
+        query_shape = ctx.get_shape(node, 0)
+        total_sequence_length = None
+        output_dtype = None
+        if query_shape is not None:
+            if len(query_shape) == 3:
+                key_shape = ctx.try_get_shape(node, 1)
+                output_shape = query_shape
+                if key_shape is not None and len(key_shape) == 3:
+                    value_shape = ctx.try_get_shape(node, 2)
+                    if value_shape is not None and len(value_shape) == 3:
+                        output_shape[2] = value_shape[2]
+                    total_sequence_length = key_shape[1]
+
+                output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+                vi = ctx.known_vi_[node.output[0]]
+                vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+            elif len(query_shape) == 5:
+                if isinstance(query_shape[2], int) and isinstance(query_shape[4], int):
+                    output_shape = [query_shape[0], query_shape[1], query_shape[2] * query_shape[4]]
+                else:
+                    output_shape = [query_shape[0], query_shape[1], f"{query_shape[2]}*{query_shape[4]}"]
+
+                total_sequence_length = query_shape[1]
+
+                output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+                vi = ctx.known_vi_[node.output[0]]
+                vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+            if len(node.output) > 1:
+                batch_size = query_shape[0]
+                num_heads = get_attribute(node, "num_heads")
+
+                head_size = None
+                if len(query_shape) == 3:
+                    head_size = (
+                        int(query_shape[2] / num_heads)
+                        if isinstance(query_shape[2], int)
+                        else f"{query_shape[2]}/{num_heads}"
+                    )
+                else:
+                    head_size = query_shape[4]
+
+                past_shape = ctx.try_get_shape(node, 6)
+
+                if past_shape is not None:
+                    if isinstance(past_shape[2], int) and isinstance(total_sequence_length, int):
+                        total_sequence_length = past_shape[2] + total_sequence_length
+                    else:
+                        total_sequence_length = f"{past_shape[2]}+{total_sequence_length}"
+
+                present_shape = [batch_size, num_heads, total_sequence_length, head_size]
+
+                assert output_dtype is not None
+                if len(node.output) > 2 and node.output[1] and node.output[2]:
+                    vi = ctx.known_vi_[node.output[1]]
+                    vi.CopyFrom(helper.make_tensor_value_info(vi.name, output_dtype, present_shape))
+                    vi = ctx.known_vi_[node.output[2]]
+                    vi.CopyFrom(helper.make_tensor_value_info(vi.name, output_dtype, present_shape))
+
+
+register_shape_handler(MultiHeadAttentionHandler())

onnxslim/core/shape_inference/contrib_ops/attention/multi_scale_deformable_attn.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for MultiScaleDeformableAttnTRT operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class MultiScaleDeformableAttnTRTHandler(ShapeHandler):
+    """Handler for MultiScaleDeformableAttnTRT operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "MultiScaleDeformableAttnTRT"
+
+    def infer_shape(self, node, ctx) -> None:
+        shape_value = ctx.try_get_shape(node, 0)
+        sampling_locations = ctx.try_get_shape(node, 3)
+        output_shape = shape_value
+        output_shape[1] = sampling_locations[1]
+        output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+
+register_shape_handler(MultiScaleDeformableAttnTRTHandler())

onnxslim/core/shape_inference/contrib_ops/attention/packed_attention.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for PackedAttention operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class PackedAttentionHandler(ShapeHandler):
+    """Handler for PackedAttention operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "PackedAttention"
+
+    def infer_shape(self, node, ctx) -> None:
+        shape = ctx.get_shape(node, 0)
+        shape_weights = ctx.get_shape(node, 1)
+        shape_bias = ctx.try_get_shape(node, 2)
+        if shape_bias is not None:
+            assert len(shape_bias) == 1
+        tripled_hidden_size = shape_bias[0] if shape_bias is not None else shape_weights[1]
+        if shape and len(shape) == 2:
+            qkv_hidden_sizes_attr = get_attribute(node, "qkv_hidden_sizes")
+            if qkv_hidden_sizes_attr is not None:
+                assert len(qkv_hidden_sizes_attr) == 3
+                shape[1] = int(qkv_hidden_sizes_attr[2])
+            elif isinstance(tripled_hidden_size, int):
+                shape[1] = int(tripled_hidden_size / 3)
+            output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+            vi = ctx.known_vi_[node.output[0]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, shape))
+
+
+register_shape_handler(PackedAttentionHandler())

onnxslim/core/shape_inference/contrib_ops/attention/packed_multi_head_attention.py

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for PackedMultiHeadAttention operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class PackedMultiHeadAttentionHandler(ShapeHandler):
+    """Handler for PackedMultiHeadAttention operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "PackedMultiHeadAttention"
+
+    def infer_shape(self, node, ctx) -> None:
+        shape_value = ctx.try_get_shape(node, 2)
+        if shape_value is not None and len(shape_value) == 2:
+            output_shape = shape_value
+        else:
+            shape_query = ctx.get_shape(node, 0)
+            assert shape_query is not None and len(shape_query) == 4
+            output_shape = [shape_query[0], shape_query[1] * shape_query[3]]
+
+        output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+
+register_shape_handler(PackedMultiHeadAttentionHandler())

onnxslim/core/shape_inference/contrib_ops/attention/remove_padding.py

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for RemovePadding operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class RemovePaddingHandler(ShapeHandler):
+    """Handler for RemovePadding operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "RemovePadding"
+
+    def infer_shape(self, node, ctx) -> None:
+        shape = ctx.get_shape(node, 0)
+        if shape and len(shape) == 3:
+            output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+            vi = ctx.known_vi_[node.output[0]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, ["token_count", shape[2]]))
+
+            vi_token_offset = ctx.known_vi_[node.output[1]]
+            vi_token_offset.CopyFrom(
+                helper.make_tensor_value_info(node.output[1], onnx.TensorProto.INT32, [shape[0], shape[1]])
+            )
+
+            vi_cumulated_seq_len = ctx.known_vi_[node.output[2]]
+            vi_cumulated_seq_len.CopyFrom(
+                helper.make_tensor_value_info(node.output[2], onnx.TensorProto.INT32, ["batch_size + 1"])
+            )
+
+            vi_max_seq_len = ctx.known_vi_[node.output[3]]
+            vi_max_seq_len.CopyFrom(helper.make_tensor_value_info(node.output[3], onnx.TensorProto.INT32, [1]))
+
+
+register_shape_handler(RemovePaddingHandler())

onnxslim/core/shape_inference/contrib_ops/attention/restore_padding.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for RestorePadding operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class RestorePaddingHandler(ShapeHandler):
+    """Handler for RestorePadding operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "RestorePadding"
+
+    def infer_shape(self, node, ctx) -> None:
+        shape_input = ctx.get_shape(node, 0)
+        shape_token_offset = ctx.get_shape(node, 1)
+        if shape_input and len(shape_input) == 2 and shape_token_offset and len(shape_token_offset) == 2:
+            output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+            vi = ctx.known_vi_[node.output[0]]
+            output_shape = [shape_token_offset[0], shape_token_offset[1], shape_input[1]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+
+register_shape_handler(RestorePaddingHandler())

onnxslim/core/shape_inference/contrib_ops/misc/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Miscellaneous contrib operator shape handlers."""
+
+from . import bias_gelu
+from . import fast_gelu
+from . import gelu
+from . import quick_gelu
+from . import gemm_fast_gelu
+from . import gemm_float8
+from . import bias_split_gelu
+from . import bias_add
+from . import rotary_embedding
+from . import python_op

onnxslim/core/shape_inference/contrib_ops/misc/bias_add.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for BiasAdd operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class BiasAddHandler(ShapeHandler):
+    """Handler for BiasAdd operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "BiasAdd"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+
+register_shape_handler(BiasAddHandler())

onnxslim/core/shape_inference/contrib_ops/misc/bias_gelu.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for BiasGelu operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class BiasGeluHandler(ShapeHandler):
+    """Handler for BiasGelu operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "BiasGelu"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+
+register_shape_handler(BiasGeluHandler())

onnxslim/core/shape_inference/contrib_ops/misc/bias_split_gelu.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for BiasSplitGelu operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class BiasSplitGeluHandler(ShapeHandler):
+    """Handler for BiasSplitGelu operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "BiasSplitGelu"
+
+    def infer_shape(self, node, ctx) -> None:
+        input_shape = ctx.get_shape(node, 0)
+        bias_shape = ctx.get_shape(node, 1)
+        if input_shape and bias_shape and isinstance(bias_shape[0], int):
+            output_shape = input_shape
+            output_shape[2] = int(bias_shape[0] / 2)
+            vi = ctx.known_vi_[node.output[0]]
+            output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+            vi.CopyFrom(helper.make_tensor_value_info(vi.name, output_dtype, output_shape))
+
+
+register_shape_handler(BiasSplitGeluHandler())

onnxslim/core/shape_inference/contrib_ops/misc/fast_gelu.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for FastGelu operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class FastGeluHandler(ShapeHandler):
+    """Handler for FastGelu operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "FastGelu"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+
+register_shape_handler(FastGeluHandler())

onnxslim/core/shape_inference/contrib_ops/misc/gelu.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Gelu operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class GeluHandler(ShapeHandler):
+    """Handler for Gelu operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Gelu"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+
+register_shape_handler(GeluHandler())

onnxslim/core/shape_inference/contrib_ops/misc/gemm_fast_gelu.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for GemmFastGelu operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class GemmFastGeluHandler(ShapeHandler):
+    """Handler for GemmFastGelu operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "GemmFastGelu"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.compute_matmul_shape(node)
+
+
+register_shape_handler(GemmFastGeluHandler())

onnxslim/core/shape_inference/contrib_ops/misc/gemm_float8.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for GemmFloat8 operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class GemmFloat8Handler(ShapeHandler):
+    """Handler for GemmFloat8 operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "GemmFloat8"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.compute_matmul_shape(node)
+
+
+register_shape_handler(GemmFloat8Handler())

onnxslim/core/shape_inference/contrib_ops/misc/python_op.py

@@ -0,0 +1,67 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for PythonOp operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_shape_from_sympy_shape
+
+
+class PythonOpHandler(ShapeHandler):
+    """Handler for PythonOp operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "PythonOp"
+
+    def infer_shape(self, node, ctx) -> None:
+        output_tensor_types = get_attribute(node, "output_tensor_types")
+        assert output_tensor_types, f"PythonOp '{node.name}' has no output_tensor_types attribute."
+        output_tensor_ranks = get_attribute(node, "output_tensor_ranks")
+        assert output_tensor_ranks, f"PythonOp '{node.name}' has no output_tensor_ranks attribute."
+
+        try:
+            from onnxruntime.capi._pybind_state import get_shape_inference_function
+
+            func_name = get_attribute(node, "func_name").decode()
+            shape_inferer = get_shape_inference_function(func_name)
+        except ImportError:
+            shape_inferer = None
+
+        # Set the context output separately.
+        # The first output is torch.autograd.Function's context.
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], onnx.TensorProto.INT64, []))
+
+        if shape_inferer is not None:
+            input_shapes = []
+            input_dtypes = []
+            for input_index in range(len(node.input)):
+                shape = ctx.get_shape(node, input_index)
+                input_shapes.append(shape)
+                input_dtype = ctx.known_vi_[node.input[input_index]].type.tensor_type.elem_type
+                input_dtypes.append(input_dtype)
+            output_shapes, output_dtypes = shape_inferer(node, input_shapes, input_dtypes)
+            assert len(output_shapes) == len(output_dtypes) == (len(node.output) - 1), (
+                f"PythonOp '{func_name}' returned {len(output_shapes)} shapes and {len(output_dtypes)} dtypes, "
+                f"but expected {len(node.output) - 1} outputs."
+            )
+            for i in range(len(node.output) - 1):
+                output_index = i + 1
+                vi = ctx.known_vi_[node.output[output_index]]
+                vi.CopyFrom(helper.make_tensor_value_info(node.output[output_index], output_dtypes[i], output_shapes[i]))
+        else:
+            # General shape inference for PythonOp.
+            for i in range(len(node.output) - 1):
+                vi = ctx.known_vi_[node.output[i + 1]]
+                sympy_shape = ctx.new_symbolic_shape(output_tensor_ranks[i], node)
+                shape = get_shape_from_sympy_shape(sympy_shape)
+                value_info = helper.make_tensor_value_info(node.output[i + 1], output_tensor_types[i], shape)
+                vi.CopyFrom(value_info)
+
+
+register_shape_handler(PythonOpHandler())