Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions onnxscript/function_libs/torch_lib/ops/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -8764,6 +8764,11 @@ def aten_type_as(self: TTensor, other: TTensor2) -> TTensor2:
def aten_unbind(self: TTensor, dim: int = 0) -> Sequence[TTensor]:
"""unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[]"""

if isinstance(self.shape[dim], int):
# We can create a definitive split op if the input shape is static
outputs = op.Split(self, axis=dim, num_outputs=self.shape[dim])
return [op.Squeeze(out, [self.shape[dim]]) for out in outputs]

return op.SplitToSequence(self, axis=dim, keepdims=False)


Expand Down
Loading