Propagate context to waiters correctly. We sometimes resume execution from a shared thread (particularly in the TPU code) and currently the context is lost to all following operations.

PiperOrigin-RevId: 754113855

Author: tensorflower-gardener

tsl/platform/BUILD

@@ -661,6 +661,7 @@ cc_library(
 filegroup(
     name = "xla_cpu_runtime_srcs",
     srcs = [
+        "context.h",
         "env_time.h",
         "macros.h",
         "platform.h",