Skip to content

Commit

Permalink
[torchlib] Update operator:pow implementation (#2069)
Browse files Browse the repository at this point in the history
Register it to aten_pow instead because exponent may not be a tensor.

Fixes pytorch/pytorch#147606
  • Loading branch information
justinchuby authored Feb 21, 2025
1 parent 013d28c commit d4bbee7
Showing 1 changed file with 8 additions and 9 deletions.
17 changes: 8 additions & 9 deletions onnxscript/function_libs/torch_lib/ops/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -6651,22 +6651,21 @@ def aten_positive(self: TensorType) -> TensorType:
raise NotImplementedError()


@torch_op(
("aten::pow.Tensor_Tensor", "aten::pow.Tensor_Scalar"),
trace_only=True,
)
@torch_op(("aten::pow.Tensor_Tensor", "_operator::pow"), trace_only=True)
def aten_pow(self: TReal, exponent: TTensor) -> TReal:
"""pow(Tensor self, Tensor exponent) -> Tensor"""
return op.Pow(self, exponent)


@torch_op(
("_operator::pow", "aten::pow.Scalar"),
trace_only=True,
)
@torch_op("aten::pow.Tensor_Scalar", trace_only=True)
def aten_pow_tensor_scalar(self: TReal, exponent: float) -> TReal:
"""pow(Tensor self, Scalar exponent) -> Tensor"""
return op.Pow(self, exponent)


@torch_op("aten::pow.Scalar", trace_only=True)
def aten_pow_scalar(self: float, exponent: TTensor) -> TTensor:
"""pow.Scalar(Scalar self, Tensor exponent) -> Tensor"""

return op.Pow(op.Cast(self, to=exponent.dtype), exponent)


Expand Down

0 comments on commit d4bbee7

Please sign in to comment.