|
@torch_op("aten::unbind.int") |
|
def aten_unbind(self: TTensor, dim: int = 0) -> Sequence[TTensor]: |
|
"""unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[]""" |
|
|
|
split_sizes = op.Constant(value_int=1) |
|
return op.SplitToSequence(self, split_sizes, axis=dim, keepdims=False) |
When split args are specified, SplitToSequence disables keepdims, resulting in tensor sequences with shapes that differ from those output by aten::unbind.int.
cf. https://onnx.ai/onnx/operators/onnx__SplitToSequence.html#attributes
onnxscript/onnxscript/function_libs/torch_lib/ops/core.py
Lines 8619 to 8624 in b2d94fe
When split args are specified,
SplitToSequencedisableskeepdims, resulting in tensor sequences with shapes that differ from those output byaten::unbind.int.cf. https://onnx.ai/onnx/operators/onnx__SplitToSequence.html#attributes