From caf22facc69a872cb81c7143218e2323cbae7205 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Jun 2024 12:03:11 -0700 Subject: [PATCH] [torchlib] Fix names for registered functions (#1606) Fix the name for `getitem` and `adaptive*` ops. --- onnxscript/function_libs/torch_lib/ops/core.py | 11 +++++------ onnxscript/function_libs/torch_lib/ops/nn.py | 6 +++--- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 292aeab01..6d15f52f7 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -3560,7 +3560,7 @@ def aten_full(size: INT64, fill_value: FLOAT, dtype: int = FLOAT.dtype): @torch_op("aten::full_like") -def aten_full_like(self, fill_value: TTensor) -> TTensor: +def aten_full_like(self: TTensor, fill_value: TTensor) -> TTensor: """full_like(Tensor self, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor""" fill_value = op.CastLike(fill_value, self) @@ -3570,7 +3570,7 @@ def aten_full_like(self, fill_value: TTensor) -> TTensor: @torch_op("aten::full_like") -def aten_full_like_dtype(self, fill_value: TTensor, dtype: int) -> TTensor: +def aten_full_like_dtype(self: TTensor, fill_value: TTensor, dtype: int) -> TTensor: """full_like(Tensor self, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor""" fill_value = op.Cast(fill_value, to=dtype) @@ -3669,8 +3669,7 @@ def aten_ger(self: TensorType, vec2: TensorType) -> TensorType: raise NotImplementedError() -# NOTE: The name is made up for `getitem` to be included in the registry -@torch_op("aten::getitem") +@torch_op("_operator::getitem") def aten_getitem(self: Sequence[TTensor], i: INT64) -> TTensor: return op.SequenceAt(self, i) @@ -8174,7 +8173,7 @@ def aten_trace_backward(grad: TensorType, sizes: INT64) -> TensorType: @torch_op(("aten::transpose", "aten::transpose.int"), trace_only=True) -def aten_transpose(self, dim0: int, dim1: int): +def aten_transpose(self: TTensor, dim0: int, dim1: int) -> TTensor: """transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a)""" # Use trace only to construct the prem attribute in Transpose @@ -8194,7 +8193,7 @@ def aten_transpose(self, dim0: int, dim1: int): @torch_op(("aten::transpose", "aten::transpose.int"), trace_only=True, complex=True) -def aten_transpose_complex(self, dim0: int, dim1: int): +def aten_transpose_complex(self: TTensor, dim0: int, dim1: int) -> TTensor: """transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a)""" # Use trace only to construct the prem attribute in Transpose diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index 46205f296..85fc4597c 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -40,7 +40,7 @@ TFloatUnlessFloat32 = TypeVar("TFloatUnlessFloat32", bound=Union[BFLOAT16, FLOAT16, DOUBLE]) -@torch_op("aten::aten_adaptive_avg_pool1d", traceable=True) +@torch_op("aten::adaptive_avg_pool1d", traceable=True) def aten_adaptive_avg_pool1d(self: TFloat, output_size: INT64[1]) -> TFloat: """adaptive_avg_pool1d(Tensor self, int[1] output_size) -> Tensor""" @@ -58,7 +58,7 @@ def aten_adaptive_avg_pool1d(self: TFloat, output_size: INT64[1]) -> TFloat: return result -@torch_op("aten::aten_adaptive_avg_pool2d", traceable=True) +@torch_op("aten::adaptive_avg_pool2d", traceable=True) def aten_adaptive_avg_pool2d(self: TFloat, output_size: INT64[2]) -> TFloat: """adaptive_avg_pool2d(Tensor self, SymInt[2] output_size) -> Tensor""" @@ -76,7 +76,7 @@ def aten_adaptive_avg_pool2d(self: TFloat, output_size: INT64[2]) -> TFloat: return result -@torch_op("aten::aten_adaptive_avg_pool3d", traceable=True) +@torch_op("aten::adaptive_avg_pool3d", traceable=True) def aten_adaptive_avg_pool3d(self: TFloat, output_size: INT64[3]) -> TFloat: """adaptive_avg_pool3d(Tensor self, SymInt[3] output_size) -> Tensor"""