Checking if the input is finite before calculation in lowering of pow func (pytorch#167723)

Fixes pytorch#167197

The inductor backend is trying to convert the float infinity value to an integer in pow lowering (possibly for indexing, iteration counts, or type conversions). Python/C++ cannot convert float('inf') to an integer, causing the overflow error

Pull Request resolved: pytorch#167723
Approved by: https://github.com/shunting314

Author: krastogi-in

test/inductor/test_torchinductor.py

@@ -5528,6 +5528,32 @@ def fn(x):
             check_lowp=not is_halide_backend(self.device),  # misaligned addr fp16
         )
 
+    def test_lp_pool1d_with_inf_norm(self):
+        # https://github.com/pytorch/pytorch/issues/167197
+        # Test that LPPool1d works with infinity norm (should behave like max pooling)
+        def fn(x):
+            return torch.nn.functional.lp_pool1d(
+                x, norm_type=float("inf"), kernel_size=2, stride=2
+            )
+
+        self.common(
+            fn,
+            (torch.randn(3, 4, 8),),
+        )
+
+    def test_lp_pool2d_with_inf_norm(self):
+        # https://github.com/pytorch/pytorch/issues/167197
+        # Test that LPPool2d works with infinity norm (should behave like max pooling)
+        def fn(x):
+            return torch.nn.functional.lp_pool2d(
+                x, norm_type=float("inf"), kernel_size=2, stride=2
+            )
+
+        self.common(
+            fn,
+            (torch.randn(3, 4, 8, 8),),
+        )
+
     @tf32_on_and_off(0.006)
     @skip_if_gpu_halide  # slow
     def test_alexnet_prefix(self):
@@ -6307,6 +6333,15 @@ def fn(x):
         x = torch.randn([16, 16], device=self.device)
         self.assertEqual(cfn(x), fn(x))
 
+    def test_pow_infinite(self):
+        def fn(a, b):
+            return torch.pow(a, b)
+
+        opt = torch.compile(fn, backend="inductor")
+        a = torch.randn((3, 4, 8), device=self.device)
+        b = float("inf")
+        self.assertTrue(same(opt(a, b), fn(a, b)))
+
     def test_glu(self):
         def fn(x):
             return aten.glu(x, -1), aten.glu(x, 1), aten.glu(x, 2)

torch/_inductor/lowering.py

@@ -6361,7 +6361,7 @@ def pow_native(a, b):
 
 @register_lowering(aten.pow, broadcast=True)
 def pow(a, b):
-    if isinstance(b, float) and b == int(b):
+    if isinstance(b, float) and math.isfinite(b) and b == int(b):
         return pow(a, int(b))
     elif isinstance(b, float) and b == 0.5:
         return sqrt(a)