Enable cross entropy loss for xla autocast with FP32 precision (#7992)

test/test_bf16_autocast.py

@@ -0,0 +1,29 @@
+import os
+import re
+import torch
+import torch_xla
+import torch_xla.core.xla_model as xm
+import unittest
+
+device = xm.xla_device()
+
+
+class TestAutocastXla(unittest.TestCase):
+
+ def test_cross_entropy_loss(self):
+ data = torch.randn(16, 10).to(torch.bfloat16).to(device)
+ target = torch.randn(16, 10).to(torch.bfloat16).to(device)
+ with torch.autocast("xla"):
+ loss = torch.nn.CrossEntropyLoss()(data, target)
+ hlo = torch_xla._XLAC._get_xla_tensors_hlo([loss])
+ self.assertTrue(
+ re.search(rf".*convert.*f32.*convert.*bf16", hlo) is not None)
+
+ self.assertTrue(
+ re.search(rf".*exponential.*f32.*exponential.*f32", hlo) is not None)
+
+ self.assertTrue(re.search(rf".*log.*f32.*log.*f32", hlo) is not None)
+
+
+if __name__ == "__main__":
+ unittest.main()

torch_xla/csrc/autocast_mode.cpp

@@ -95,7 +95,7 @@ TORCH_LIBRARY_IMPL(aten, AutocastXLA, m) {
  KERNEL_XLA(hinge_embedding_loss, fp32)
  // KERNEL_XLA(poisson_nll_loss, fp32)
  KERNEL_XLA(smooth_l1_loss, fp32)
- // KERNEL_XLA(cross_entropy_loss, fp32)
+ KERNEL_XLA(cross_entropy_loss, fp32)
  KERNEL_XLA(l1_loss, fp32)
  // KERNEL_XLA(huber_loss, fp32)
  KERNEL_XLA(margin_ranking_loss, fp32)