From 96e2e31af942e5b30d6bfd171f894c0ec7788651 Mon Sep 17 00:00:00 2001 From: xadupre Date: Thu, 3 Jul 2025 11:07:25 +0200 Subject: [PATCH 1/3] Improves aten_chunk conversion --- .../index.md | 3 +++ .../function_libs/torch_lib/ops/core.py | 25 +++---------------- 2 files changed, 7 insertions(+), 21 deletions(-) diff --git a/docs/articles/2023-06-20-introducing-onnx-script/index.md b/docs/articles/2023-06-20-introducing-onnx-script/index.md index bf238af983..d9b230b197 100644 --- a/docs/articles/2023-06-20-introducing-onnx-script/index.md +++ b/docs/articles/2023-06-20-introducing-onnx-script/index.md @@ -187,8 +187,11 @@ We will cover the new PyTorch ONNX exporter in a separate post with more depth a By deeply weaving ONNX Script support into the PyTorch ONNX exporter, we have also made it possible to augment PyTorch model code with specialized ONNX functions as custom operators. We introduced initial support for this in the TorchScript exporter starting with [PyTorch 1.13][torch-onnx-customops] and continue to refine this capability in the new exporter. ## An End-to-End Example + Let’s look at an example slightly more complicated than GELU. In fact, the following example is adapted directly from the new PyTorch ONNX exporter, implementing support for [`torch.chunk`][torch-chunk], which attempts to split a tensor into the number of specified chunks. +.. TODO: change this example + ```python from typing import Sequence from onnxscript import opset18 as op, script, FLOAT, INT64 diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 92b8abb36d..85c0b168e0 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -1647,29 +1647,12 @@ def aten_choose_qparams_optimized( raise NotImplementedError() -@torch_op("aten::chunk") +@torch_op("aten::chunk", trace_only=True) def aten_chunk(self: TTensor, chunks: int, dim: int = 0) -> Sequence[TTensor]: """chunk(Tensor(a -> *) self, int chunks, int dim=0) -> Tensor(a)[]""" - # This will create a Sequence of tensors - neg_1 = op.Constant(value_ints=[-1]) - # Get size of specified dim - self_shape = op.Shape(self) - dim_size = op.Gather(self_shape, dim, axis=0) - # Compute size/chunk to get the number of data in one chunk - num_per_chunk = op.Div(dim_size, chunks) - num_per_chunk = op.Cast(op.Mod(dim_size, chunks) > 0, to=INT64.dtype) + num_per_chunk # type: ignore[operator] - - # Compute real chunk number - num_chunk = op.Div(dim_size, num_per_chunk) - # Get something like [n, n, n, n, ...], total num_chunk - list_split = op.Expand(num_per_chunk, op.Reshape(num_chunk, neg_1)) - - remainder = op.Mod(dim_size, num_per_chunk) - if remainder > 0: # type: ignore[operator] - # Append the remainder to the [n, n, n, n, ..., r] - list_split = op.Concat(list_split, op.Reshape(remainder, neg_1), axis=0) - - return op.SplitToSequence(self, list_split, axis=dim) + if chunks == 1: + return g.op.Identity(self) + return g.op.Split(self, axis=dim, num_outputs=chunks) @torch_op("aten::clamp", trace_only=True) From e7f21f71e6daab92b2bdf732200bcef8ffdec799 Mon Sep 17 00:00:00 2001 From: xadupre Date: Tue, 22 Jul 2025 09:37:37 +0200 Subject: [PATCH 2/3] removes g.op --- onnxscript/function_libs/torch_lib/ops/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 85c0b168e0..da16bba726 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -1651,8 +1651,8 @@ def aten_choose_qparams_optimized( def aten_chunk(self: TTensor, chunks: int, dim: int = 0) -> Sequence[TTensor]: """chunk(Tensor(a -> *) self, int chunks, int dim=0) -> Tensor(a)[]""" if chunks == 1: - return g.op.Identity(self) - return g.op.Split(self, axis=dim, num_outputs=chunks) + return op.Identity(self) + return op.Split(self, axis=dim, num_outputs=chunks) @torch_op("aten::clamp", trace_only=True) From ae1a0d6d904caa0feed25ac28077559aa6c7fa4a Mon Sep 17 00:00:00 2001 From: xadupre Date: Tue, 22 Jul 2025 09:38:56 +0200 Subject: [PATCH 3/3] remove --- docs/articles/2023-06-20-introducing-onnx-script/index.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/docs/articles/2023-06-20-introducing-onnx-script/index.md b/docs/articles/2023-06-20-introducing-onnx-script/index.md index d9b230b197..bf238af983 100644 --- a/docs/articles/2023-06-20-introducing-onnx-script/index.md +++ b/docs/articles/2023-06-20-introducing-onnx-script/index.md @@ -187,11 +187,8 @@ We will cover the new PyTorch ONNX exporter in a separate post with more depth a By deeply weaving ONNX Script support into the PyTorch ONNX exporter, we have also made it possible to augment PyTorch model code with specialized ONNX functions as custom operators. We introduced initial support for this in the TorchScript exporter starting with [PyTorch 1.13][torch-onnx-customops] and continue to refine this capability in the new exporter. ## An End-to-End Example - Let’s look at an example slightly more complicated than GELU. In fact, the following example is adapted directly from the new PyTorch ONNX exporter, implementing support for [`torch.chunk`][torch-chunk], which attempts to split a tensor into the number of specified chunks. -.. TODO: change this example - ```python from typing import Sequence from onnxscript import opset18 as op, script, FLOAT, INT64