Revert the mul change (#7271) (#7285)

pytorch · Jun 17, 2024 · 7296133 · 7296133
1 parent fda5828
commit 7296133
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/test/test_operations.py b/test/test_operations.py
@@ -2099,6 +2099,8 @@ def test(f, xshape, ishapes):
     for xshape, i0shape, i1shape in cases[f2]:
       test(f2, xshape, (i0shape, i1shape))
 
+  @unittest.skipIf(
+      True, "skip since https://github.com/pytorch/xla/pull/7130 is reverted")
   def test_inplace_mul_scalar_different_dtype(self):
     # This tests whether the returned output data-type agrees on PyTorch
     # and XLA sides.

diff --git a/torch_xla/csrc/aten_xla_type.cpp b/torch_xla/csrc/aten_xla_type.cpp
@@ -2221,14 +2221,11 @@ at::Tensor XLANativeFunctions::mse_loss_backward(const at::Tensor& grad_output,
 at::Tensor XLANativeFunctions::mul(const at::Tensor& self,
                                    const at::Tensor& other) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
-  using FnType = XLATensorPtr(const XLATensorPtr&, const XLATensorPtr&,
-                              std::optional<at::ScalarType>);
-  return OpConfig::From(static_cast<FnType*>(tensor_methods::mul))
-      .add_input(self)
-      .add_input(other)
-      .cast_inputs_to_common_dtype()
-      .use_opmathtype_for_compute()
-      .run();
+  return DoBinaryOp(self, other,
+                    [&](const XLATensorPtr& xself, const XLATensorPtr& xother,
+                        at::ScalarType dtype) {
+                      return tensor_methods::mul(xself, xother, dtype);
+                    });
 }
 
 at::Tensor XLANativeFunctions::mul(const at::Tensor& self,