From 62a5797e7ebd94be044d6eaa9edf3e188d703100 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Sun, 9 Nov 2025 18:37:24 +0100 Subject: [PATCH 1/5] Add converter for unique_consecutive --- .../function_libs/torch_lib/ops/core.py | 47 ++++++++++++++++++- .../function_libs/torch_lib/e2e_ops_tests.py | 30 ++++++++++++ 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 2cbecdcfc2..566544b220 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -51,6 +51,7 @@ from onnxscript.onnx_opset import opset18 as op from onnxscript.onnx_types import TensorType +_INT32_MAX = 2147483647 _INT64_MAX = 9223372036854775807 _INT64_MIN = -9223372036854775808 _MATH_PI = math.pi @@ -9086,15 +9087,57 @@ def aten_unfold_copy(self: TensorType, dimension: int, size: int, step: int) -> raise NotImplementedError() +@torch_op("aten::unique_consecutive", trace_only=True) def aten_unique_consecutive( - self: TensorType, + x: TensorType, return_inverse: bool = False, return_counts: bool = False, dim: Optional[int] = None, ) -> tuple[TensorType, TensorType, TensorType]: """unique_consecutive(Tensor self, bool return_inverse=False, bool return_counts=False, int? dim=None) -> (Tensor, Tensor, Tensor)""" + assert x.dtype == INT64.dtype or x.dtype == INT32.dtype, ( + "unique_consecutive not implemented for other type than int32, int64" + ) + rank_x = len(x.shape) - raise NotImplementedError() + zero = op.Constant(value=ir.tensor([0], dtype=x.dtype)) + minus_one = op.Constant(value=ir.tensor([-1], dtype=x.dtype)) + + if dim is None: + if rank_x != 1: + x = op.Reshape(x, minus_one) + dim = 0 + else: + assert rank_x == 1 and dim == 0, ( + f"Not implemented for x={x!r} with rank={rank_x} and dim={dim}." + ) + + lag = op.Concat( + # Hopefully this will never be equal to the first value of the tensor x + # ideally we could do differently but with a higher cost + op.Constant(value=ir.tensor([_INT32_MAX], dtype=x.dtype)), + op.Slice(x, zero, minus_one, zero), + axis=0, + ) + eq = op.Equal(x, lag) + diff = op.Not(eq) + res = op.Compress(x, diff, axis=0) + + zero_no_dim = op.Constant(value=ir.tensor(0, dtype=x.dtype)) + one_no_dim = op.Constant(value=ir.tensor(1, dtype=x.dtype)) + one = op.Constant(value=ir.tensor([1], dtype=x.dtype)) + + inverse = op.Sub(op.CumSum(op.Cast(diff, to=x.dtype), zero), one) + shape_x = op.Shape(x) + indices = op.Range(zero_no_dim, op.Squeeze(shape_x), one_no_dim) + points = op.Compress(indices, diff, axis=0) + lagp = op.Concat( + op.Slice(points, one, op.Shape(points), zero), + shape_x, + axis=0, + ) + counts = op.Sub(lagp, points) + return res, inverse, counts @torch_op("aten::_unique", trace_only=True) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index f74dda699d..1fded13d1c 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -406,6 +406,36 @@ def forward(self, x): onnx_program = torch.onnx.export(model, (x,), dynamo=True, verbose=False) _testing.assert_onnx_program(onnx_program) + def test_aten_unique_consecutive(self): + class Model(torch.nn.Module): + def forward(self, x): + return torch.unique_consecutive(x) + + model = Model() + x = torch.tensor([0, 1, 2, 2, 3, 3, 0, 0], dtype=torch.int64) + onnx_program = torch.onnx.export( + model, + (x,), + dynamic_shapes=({0: "length"},), + dynamo=True, + ) + _testing.assert_onnx_program(onnx_program) + + def test_aten_unique_consecutive_return(self): + class Model(torch.nn.Module): + def forward(self, x): + return torch.unique_consecutive(x, return_inverse=True, return_counts=True) + + model = Model() + x = torch.tensor([0, 1, 2, 2, 3, 3, 3, 0, 0], dtype=torch.int64) + onnx_program = torch.onnx.export( + model, + (x,), + dynamic_shapes=({0: "length"},), + dynamo=True, + ) + _testing.assert_onnx_program(onnx_program) + if __name__ == "__main__": unittest.main() From 1807619bd0888dbf338e8aff63bff3572d20bae7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Sun, 9 Nov 2025 18:48:22 +0100 Subject: [PATCH 2/5] Potential fix for pull request finding 'Unused local variable' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- onnxscript/function_libs/torch_lib/ops/core.py | 1 - 1 file changed, 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 566544b220..d0d5601c53 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -9106,7 +9106,6 @@ def aten_unique_consecutive( if dim is None: if rank_x != 1: x = op.Reshape(x, minus_one) - dim = 0 else: assert rank_x == 1 and dim == 0, ( f"Not implemented for x={x!r} with rank={rank_x} and dim={dim}." From fff2de7e4429acf5a9b85c361c59db65cd4683ea Mon Sep 17 00:00:00 2001 From: "G. Ramalingam" Date: Sun, 9 Nov 2025 19:50:25 -0800 Subject: [PATCH 3/5] Potential fix for code scanning alert no. 20975: Syntax error Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- tests/function_libs/torch_lib/e2e_ops_tests.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index cc2d154ef9..4b7c1482ab 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -433,6 +433,8 @@ def forward(self, x): (x,), dynamic_shapes=({0: "length"},), dynamo=True, + ) + _testing.assert_onnx_program(onnx_program) def test_aten_stft_1(self): class Model(torch.nn.Module): From 5c181fdeeaa3d7f6b5b5579475ba711e2f56d5bc Mon Sep 17 00:00:00 2001 From: "G. Ramalingam" Date: Sun, 9 Nov 2025 20:03:29 -0800 Subject: [PATCH 4/5] Address lint checker warning --- onnxscript/function_libs/torch_lib/ops/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 0c3a0072a3..a7f82ea67a 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -9192,7 +9192,7 @@ def aten_unique_consecutive( dim: Optional[int] = None, ) -> tuple[TensorType, TensorType, TensorType]: """unique_consecutive(Tensor self, bool return_inverse=False, bool return_counts=False, int? dim=None) -> (Tensor, Tensor, Tensor)""" - assert x.dtype == INT64.dtype or x.dtype == INT32.dtype, ( + assert x.dtype in {INT64.dtype, INT32.dtype}, ( "unique_consecutive not implemented for other type than int32, int64" ) rank_x = len(x.shape) From ecb30f974bd8aba1d65ad84be4cdabab0194f93d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Mon, 10 Nov 2025 09:54:41 +0100 Subject: [PATCH 5/5] fix int64 --- onnxscript/function_libs/torch_lib/ops/core.py | 5 +++-- tests/function_libs/torch_lib/e2e_ops_tests.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index a7f82ea67a..a25015b232 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -9198,7 +9198,8 @@ def aten_unique_consecutive( rank_x = len(x.shape) zero = op.Constant(value=ir.tensor([0], dtype=x.dtype)) - minus_one = op.Constant(value=ir.tensor([-1], dtype=x.dtype)) + zero64 = op.Constant(value=ir.tensor([0], dtype=INT64.dtype)) + minus_one = op.Constant(value=ir.tensor([-1], dtype=INT64.dtype)) if dim is None: if rank_x != 1: @@ -9212,7 +9213,7 @@ def aten_unique_consecutive( # Hopefully this will never be equal to the first value of the tensor x # ideally we could do differently but with a higher cost op.Constant(value=ir.tensor([_INT32_MAX], dtype=x.dtype)), - op.Slice(x, zero, minus_one, zero), + op.Slice(x, zero64, minus_one, zero64), axis=0, ) eq = op.Equal(x, lag) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 4b7c1482ab..1546de59bd 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -421,6 +421,21 @@ def forward(self, x): ) _testing.assert_onnx_program(onnx_program) + def test_aten_unique_consecutive_int32(self): + class Model(torch.nn.Module): + def forward(self, x): + return torch.unique_consecutive(x) + + model = Model() + x = torch.tensor([0, 1, 2, 2, 3, 3, 0, 0], dtype=torch.int32) + onnx_program = torch.onnx.export( + model, + (x,), + dynamic_shapes=({0: "length"},), + dynamo=True, + ) + _testing.assert_onnx_program(onnx_program) + def test_aten_unique_consecutive_return(self): class Model(torch.nn.Module): def forward(self, x):