enable XLA around model training step via jit_compile

PiperOrigin-RevId: 366330392

official/core/base_trainer.py

@@ -379,7 +379,11 @@ class Trainer(_AsyncTrainer):
     """See base class."""
 
     def step_fn(inputs):
-      logs = self.task.train_step(
+      if self.config.runtime.enable_xla and (self.config.runtime.num_gpus > 0):
+        task_train_step = tf.function(self.task.train_step, jit_compile=True)
+      else:
+        task_train_step = self.task.train_step
+      logs = task_train_step(
           inputs,
           model=self.model,
           optimizer=self.optimizer,