diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index bae2e2cfb7..77163d65e8 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -6832,6 +6832,14 @@ def aten_mul_complex(self: TReal, other: TReal) -> TReal: return op.Concat(real, imag, axis=-1) +@torch_op(("aten::mul.Scalar", "aten::multiply.Scalar"), trace_only=True) +def aten_mul_scalar(self: TTensor, other: float) -> TTensor: + """mul.Scalar(Tensor self, Scalar other) -> Tensor""" + + other = op.Constant(value=ir.tensor(other, dtype=self.dtype)) + return aten_mul(self, other) + + @torch_op("aten::multinomial", trace_only=True) def aten_multinomial( self: TFloat, diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 1bea819f11..8146510cc4 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -84,6 +84,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) _testing.assert_onnx_program(onnx_program) + def test_mul_tensor_scalar_float(self): + class Model(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) * 1.0 + + onnx_program = torch.onnx.export( + Model(), + (torch.tensor([1, 2, 3], dtype=torch.float16),), + dynamo=True, + optimize=False, + ) + _testing.assert_onnx_program(onnx_program) + def test_bincount(self): class Model(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: