Skip to content

Commit

Permalink
Add Op(aten::_to_copy) | feat(torchlib) (#1148)
Browse files Browse the repository at this point in the history
  • Loading branch information
titaiwangms authored Nov 13, 2023
1 parent ad6f9e1 commit 445343b
Showing 1 changed file with 15 additions and 1 deletion.
16 changes: 15 additions & 1 deletion onnxscript/function_libs/torch_lib/ops/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -2053,7 +2053,7 @@ def aten_convolution_overrideable(
raise NotImplementedError()


@torch_op(("aten::copy", "aten::_to_copy"))
@torch_op("aten::copy")
def aten_copy(
self: TTensor, src: TTensor, non_blocking: bool = False # pylint: disable=unused-argument
) -> TTensor:
Expand All @@ -2063,6 +2063,20 @@ def aten_copy(
return self


@torch_op("aten::_to_copy", trace_only=True)
def aten__to_copy(
self: TTensor,
dtype: int = -1,
non_blocking: bool = False, # pylint: disable=unused-argument
) -> TTensor:
"""_to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor"""

if dtype == -1:
return op.Identity(self)
else:
return common_ops.cast_to(self, dtype=dtype)


def aten_copysign(self: TensorType, other: TensorType) -> TensorType:
"""copysign.Tensor(Tensor self, Tensor other) -> Tensor"""

Expand Down

0 comments on commit 445343b

Please sign in to comment.