From e79bbed8a2a71d0f5d1ecd26e746f39487fed4b2 Mon Sep 17 00:00:00 2001
From: DeepMind Team <no-reply@deepmind.com>
Date: Tue, 29 Aug 2023 09:24:51 -0700
Subject: [PATCH] Internal

PiperOrigin-RevId: 561052201
Change-Id: I60fd51bb8eb50f795a2d79ebdcb7c015f3ba2a0e
GitOrigin-RevId: 820eccc7d8d0994b3420b2fa5a1a97776a81e874
---
 xmanager/xm/core.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/xmanager/xm/core.py b/xmanager/xm/core.py
index 92f8b85..a27ae49 100644
--- a/xmanager/xm/core.py
+++ b/xmanager/xm/core.py
@@ -353,6 +353,13 @@ def add(
       An awaitable that would be fulfilled when the job is launched.
     """
     # pyformat: enable
+
+    # Prioritize the identity given directly to the work unit at work unit
+    # creation time, as opposed to the identity passed when adding jobs to it as
+    # this is more consistent between job generator work units and regular work
+    # units.
+    identity = self.identity or identity
+
     job = job_operators.shallow_copy_job_type(job)
     if args is not None:
       _apply_args(job, args)