Skip to content

Commit

Permalink
[torchlib] Do not register rsub (#1907)
Browse files Browse the repository at this point in the history
Remove rsub since it is handled by decomp, and torch doesn't have a type
promotion rule for rsub so we use sub instead.

Tested with 

```python
import torch


class Model(torch.nn.Module):
    def forward(self, x):
        return 1 - x


ep = torch.export.export(Model(), (torch.tensor(1),))
print(ep)

program = torch.onnx.export(Model(), (torch.tensor(1),), dynamo=True)
print(program)

```
  • Loading branch information
justinchuby authored Oct 15, 2024
1 parent a4f3bcb commit 1544ee1
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 11 deletions.
11 changes: 2 additions & 9 deletions onnxscript/function_libs/torch_lib/ops/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -7352,18 +7352,11 @@ def aten_rsqrt(self: TFloat) -> TFloat:
return op.Reciprocal(op.Sqrt(self))


@torch_op(("aten::rsub.Tensor", "aten::rsub.Scalar"))
# Do not register rsub. It will be decomposed and type promoted by torch
def aten_rsub(self: TReal, other: TReal, alpha: float = 1.0) -> TReal:
"""rsub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor"""

return op.Sub(other, op.Mul(self, alpha))


@torch_op(("aten::rsub.Tensor", "aten::rsub.Scalar"), trace_only=True, complex=True)
def aten_rsub_complex(self: TReal, other: TReal, alpha: float = 1.0) -> TReal:
"""rsub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor"""

return aten_rsub(self, other, alpha)
raise NotImplementedError


@torch_op("aten::scalar_tensor", trace_only=True)
Expand Down
2 changes: 0 additions & 2 deletions tests/function_libs/torch_lib/ops_test_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -1360,8 +1360,6 @@ def _where_input_wrangler(
),
TorchLibOpInfo("round_decimals", core_ops.aten_round_decimals),
TorchLibOpInfo("rsqrt", core_ops.aten_rsqrt),
TorchLibOpInfo("rsub", core_ops.aten_rsub),
TorchLibOpInfo("rsub", core_ops.aten_rsub_complex, complex=True),
TorchLibOpInfo(
"scalar_tensor",
core_ops.aten_scalar_tensor,
Expand Down

0 comments on commit 1544ee1

Please sign in to comment.