Added pow to autocast policy and unit test (#8684)

Author: avizon-aws

test/test_autocast_xla.py

@@ -93,6 +93,16 @@ def test_einsum(self):
       self.assertRegex(hlo, r".*dot.*bf16")
       self.assertNotRegex(hlo, r".*dot.*f32")
 
+    def test_pow(self):
+      data = torch.randn(16, 20).to(torch.bfloat16).to(device)
+
+      with torch.autocast("xla"):
+        output = data.pow(2)
+        hlo = torch_xla._XLAC._get_xla_tensors_hlo([output])
+
+        self.assertRegex(hlo, r".*convert.*f32.*convert.*bf16")
+        self.assertRegex(hlo, r".*power.*f32.*power.*f32")
+
 
 if __name__ == "__main__":
   unittest.main()

torch_xla/csrc/autocast_mode.cpp

@@ -65,6 +65,7 @@ TORCH_LIBRARY_IMPL(aten, AutocastXLA, m) {
   KERNEL_XLA(binary_cross_entropy, fp32)
   // KERNEL_XLA(grid_sampler, fp32)
   // KERNEL_XLA(polar, fp32)
+  KERNEL_XLA2(pow, Tensor_Scalar, fp32)
   KERNEL_XLA(prod, fp32)
   KERNEL_XLA2(prod, dim_int, fp32)
   KERNEL_XLA2(prod, dim_Dimname, fp32)