Ensure that tf.summary.write() keeps metadata tensor on CPU

Previously, when in a GPU device context, the logic attempted to create the default empty metadata tensor on the GPU, which fails since it's a string tensor.

PiperOrigin-RevId: 235135723

tensorflow/python/kernel_tests/BUILD

@@ -1084,7 +1084,7 @@ tf_py_test(
     ],
 )
 
-tf_py_test(
+cuda_py_test(
     name = "summary_ops_test",
     size = "small",
     srcs = ["summary_ops_test.py"],
@@ -1107,6 +1107,7 @@ tf_py_test(
         "//tensorflow/python/keras:engine",
         "//tensorflow/python/keras:layers",
     ],
+    xla_enable_strict_auto_jit = True,
 )
 
 tf_py_test(

tensorflow/python/kernel_tests/summary_ops_test.py

@@ -144,6 +144,22 @@ class SummaryOpsCoreTest(test_util.TensorFlowTestCase):
     value = events[1].summary.value[0]
     self.assertAllEqual([b'foo', b'bar'], to_numpy(value))
 
+  @test_util.run_gpu_only
+  def testWrite_gpuDeviceContext(self):
+    logdir = self.get_temp_dir()
+    with context.eager_mode():
+      with summary_ops.create_file_writer(logdir).as_default():
+        with ops.device('/GPU:0'):
+          value = constant_op.constant(42.0)
+          step = constant_op.constant(12, dtype=dtypes.int64)
+          summary_ops.write('tag', value, step=step).numpy()
+    empty_metadata = summary_pb2.SummaryMetadata()
+    events = events_from_logdir(logdir)
+    self.assertEqual(2, len(events))
+    self.assertEqual(12, events[1].step)
+    self.assertEqual(42, to_numpy(events[1].summary.value[0]))
+    self.assertEqual(empty_metadata, events[1].summary.value[0].metadata)
+
   @test_util.also_run_as_tf_function
   def testWrite_noDefaultWriter(self):
     with context.eager_mode():

tensorflow/python/ops/summary_ops_v2.py

@@ -448,9 +448,9 @@ def write(tag, tensor, step, metadata=None, name=None):
     if context.context().summary_writer_resource is None:
       return constant_op.constant(False)
     if metadata is None:
-      serialized_metadata = constant_op.constant(b"")
+      serialized_metadata = b""
     elif hasattr(metadata, "SerializeToString"):
-      serialized_metadata = constant_op.constant(metadata.SerializeToString())
+      serialized_metadata = metadata.SerializeToString()
     else:
       serialized_metadata = metadata