Skip to content

Commit 5ba4fc7

Browse files
authored
feat(atenlib): aten_expand supports -1 dim (#483)
1 parent b0d4f24 commit 5ba4fc7

File tree

1 file changed

+2
-0
lines changed
  • onnxscript/function_libs/torch_aten/ops

1 file changed

+2
-0
lines changed

onnxscript/function_libs/torch_aten/ops/core.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2039,6 +2039,8 @@ def aten_expand(self: TTensor, size: TInt) -> TTensor:
20392039
"""expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a)"""
20402040

20412041
size = op.Cast(size, to=INT64.dtype)
2042+
# To support -1 dim.
2043+
size = op.Abs(size)
20422044
return op.Expand(self, size)
20432045

20442046

0 commit comments

Comments
 (0)