Don't fallback for pow

pytorch · Nov 28, 2023 · 2af9b6e · 2af9b6e
1 parent 3fa035e
commit 2af9b6e
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 20 deletions.
diff --git a/test/test_core_aten_ops.py b/test/test_core_aten_ops.py
@@ -3295,6 +3295,29 @@ def test_aten_pow_Tensor_Scalar_2(self):
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.pow.Tensor_Scalar, args, kwargs)
 
+  def test_aten_pow_Scalar_1(self):
+    args = (10000,
+            torch.Tensor([
+                0.0000, 0.0000, 0.0156, 0.0156, 0.0312, 0.0312, 0.0469, 0.0469,
+                0.0625, 0.0625, 0.0781, 0.0781, 0.0938, 0.0938, 0.1094, 0.1094,
+                0.1250, 0.1250, 0.1406, 0.1406, 0.1562, 0.1562, 0.1719, 0.1719,
+                0.1875, 0.1875, 0.2031, 0.2031, 0.2188, 0.2188, 0.2344, 0.2344,
+                0.2500, 0.2500, 0.2656, 0.2656, 0.2812, 0.2812, 0.2969, 0.2969,
+                0.3125, 0.3125, 0.3281, 0.3281, 0.3438, 0.3438, 0.3594, 0.3594,
+                0.3750, 0.3750, 0.3906, 0.3906, 0.4062, 0.4062, 0.4219, 0.4219,
+                0.4375, 0.4375, 0.4531, 0.4531, 0.4688, 0.4688, 0.4844, 0.4844,
+                0.5000, 0.5000, 0.5156, 0.5156, 0.5312, 0.5312, 0.5469, 0.5469,
+                0.5625, 0.5625, 0.5781, 0.5781, 0.5938, 0.5938, 0.6094, 0.6094,
+                0.6250, 0.6250, 0.6406, 0.6406, 0.6562, 0.6562, 0.6719, 0.6719,
+                0.6875, 0.6875, 0.7031, 0.7031, 0.7188, 0.7188, 0.7344, 0.7344,
+                0.7500, 0.7500, 0.7656, 0.7656, 0.7812, 0.7812, 0.7969, 0.7969,
+                0.8125, 0.8125, 0.8281, 0.8281, 0.8438, 0.8438, 0.8594, 0.8594,
+                0.8750, 0.8750, 0.8906, 0.8906, 0.9062, 0.9062, 0.9219, 0.9219,
+                0.9375, 0.9375, 0.9531, 0.9531, 0.9688, 0.9688, 0.9844, 0.9844
+            ]))
+    kwargs = dict()
+    run_export_and_compare(self, torch.ops.aten.pow.Scalar, args, kwargs)
+
   @unittest.skip
   def test_aten_pow_Tensor_Tensor_0(self):
     args = (
@@ -3313,11 +3336,10 @@ def test_aten_pow_Tensor_Tensor_1(self):
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.pow.Tensor_Tensor, args, kwargs)
 
-  @unittest.skip
   def test_aten_pow_Tensor_Tensor_2(self):
     args = (
-        torch.randint(0, 10, (10, 10)).to(torch.int32),
-        torch.randint(0, 10, (10, 10)).to(torch.int32),
+        torch.randint(0, 5, (10, 10)).to(torch.int32),
+        torch.randint(0, 5, (10, 10)).to(torch.int32),
     )
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.pow.Tensor_Tensor, args, kwargs)

diff --git a/test/test_ops.py b/test/test_ops.py
@@ -148,7 +148,6 @@ def __new__(cls, name, variant_test_name=""):
             AllowedOpInfoEntry('outer'),
             AllowedOpInfoEntry('ormqr'),
             AllowedOpInfoEntry('permute'),
-            AllowedOpInfoEntry('pow'),
             AllowedOpInfoEntry('float_power'),
             AllowedOpInfoEntry('rad2deg'),
             AllowedOpInfoEntry('real'),
@@ -347,6 +346,7 @@ def __new__(cls, name, variant_test_name=""):
             # Worked locally (but failing on CI both CPU and CUDA)
             # app.circleci.com/pipelines/github/pytorch/xla/9130/workflows/71c74f3d-1735-4328-81b5-784d6e6744da/jobs/17998
             # AllowedOpInfoEntry('var_mean'),
+            # AllowedOpInfoEntry('pow'), # for int64 don't work, likely overflowed
         }))
 
 

diff --git a/torch_xla/csrc/aten_xla_type.cpp b/torch_xla/csrc/aten_xla_type.cpp
@@ -2305,36 +2305,20 @@ at::Tensor XLANativeFunctions::permute_copy(const at::Tensor& self,
 at::Tensor XLANativeFunctions::pow(const at::Tensor& self,
                                    const at::Scalar& exponent) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
-  // xla::Pow() doesn't support integer types.
-  if (!at::native::is_floating_point(self)) {
-    return at::native::call_fallback_fn<
-        &xla_cpu_fallback, ATEN_OP2(pow, Tensor_Scalar)>::call(self, exponent);
-  }
   return bridge::AtenFromXlaTensor(
       tensor_methods::pow(bridge::GetXlaTensor(self), exponent));
 }
 
 at::Tensor XLANativeFunctions::pow(const at::Tensor& self,
                                    const at::Tensor& exponent) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
-  // xla::Pow() doesn't support integer types.
-  if (!at::native::is_floating_point(self)) {
-    return at::native::call_fallback_fn<
-        &xla_cpu_fallback, ATEN_OP2(pow, Tensor_Tensor)>::call(self, exponent);
-  }
   return bridge::AtenFromXlaTensor(tensor_methods::pow(
       bridge::GetXlaTensor(self), bridge::GetXlaTensor(exponent)));
 }
 
 at::Tensor XLANativeFunctions::pow(const at::Scalar& self,
                                    const at::Tensor& exponent) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
-  // xla::Pow() doesn't support integer types.
-  if (!self.isFloatingPoint()) {
-    return at::native::call_fallback_fn<&xla_cpu_fallback,
-                                        ATEN_OP2(pow, Scalar)>::call(self,
-                                                                     exponent);
-  }
   return bridge::AtenFromXlaTensor(
       tensor_methods::pow(self, bridge::GetXlaTensor(exponent)));
 }