Invoke TensorBoard histogram v2 in tf.compat.v1.summary.histogram.

PiperOrigin-RevId: 426010859
Change-Id: Ib8cbb24b4270dec06fe1d6b36854460398784d4d

Author: yatbear

tensorflow/python/summary/summary.py

@@ -239,6 +239,7 @@ def image(name, tensor, max_outputs=3, collections=None, family=None):
     # `tf.compat.v1.summary.merge()` API.
     return _constant_op.constant(b'')
 
+  # Fall back to legacy v1 image implementation.
   if _distribute_summary_op_util.skip_summary():
     return _constant_op.constant('')
   with _summary_op_util.summary_scope(
@@ -280,8 +281,19 @@ def histogram(name, values, collections=None, family=None):
     buffer.
 
   @compatibility(TF2)
-  This API is not compatible with eager execution and `tf.function`. To migrate
-  to TF2, please use `tf.summary.histogram` instead. Please check
+  For compatibility purposes, when invoked in TF2 where the outermost context is
+  eager mode, this API will check if there is a suitable TF2 summary writer
+  context available, and if so will forward this call to that writer instead. A
+  "suitable" writer context means that the writer is set as the default writer,
+  and there is an associated non-empty value for `step` (see
+  `tf.summary.SummaryWriter.as_default`, or alternatively
+  `tf.summary.experimental.set_step`). For the forwarded call, the arguments
+  here will be passed to the TF2 implementation of `tf.summary.histogram`, and
+  the return value will be an empty bytestring tensor, to avoid duplicate
+  summary writing. This forwarding is best-effort and not all arguments will be
+  preserved.
+
+  To migrate to TF2, please use `tf.summary.histogram` instead. Please check
   [Migrating tf.summary usage to
   TF 2.0](https://www.tensorflow.org/tensorboard/migrate#in_tf_1x) for concrete
   steps for migration.
@@ -306,6 +318,18 @@ def histogram(name, values, collections=None, family=None):
 
   @end_compatibility
   """
+  # Special case: invoke v2 op for TF2 users who have a v2 writer.
+  if _should_invoke_v2_op():
+    # Defer the import to happen inside the symbol to prevent breakage due to
+    # missing dependency.
+    from tensorboard.summary.v2 import histogram as histogram_v2  # pylint: disable=g-import-not-at-top
+    with _compat_summary_scope(name, family) as tag:
+      histogram_v2(name=tag, data=values, step=_get_step_for_v2())
+    # Return an empty Tensor, which will be acceptable as an input to the
+    # `tf.compat.v1.summary.merge()` API.
+    return _constant_op.constant(b'')
+
+  # Fall back to legacy v1 histogram implementation.
   if _distribute_summary_op_util.skip_summary():
     return _constant_op.constant('')
   with _summary_op_util.summary_scope(

tensorflow/python/summary/summary_v2_test.py

@@ -36,8 +36,8 @@ def test_scalar_summary_v2__w_writer(self):
     """Tests scalar v2 invocation with a v2 writer."""
     with test.mock.patch.object(
         summary_v2, 'scalar', autospec=True) as mock_scalar_v2:
-      with summary_ops_v2.create_summary_file_writer('/tmp/test').as_default(
-          step=1):
+      with summary_ops_v2.create_summary_file_writer(
+          self.get_temp_dir()).as_default(step=1):
         i = constant_op.constant(2.5)
         tensor = summary_lib.scalar('float', i)
     # Returns empty string.
@@ -62,7 +62,7 @@ def test_scalar_summary_v2__global_step_not_set(self):
       with test.mock.patch.object(
           summary_v2, 'scalar', autospec=True) as mock_scalar_v2:
         with summary_ops_v2.create_summary_file_writer(
-            '/tmp/test').as_default():
+            self.get_temp_dir()).as_default():
           summary_lib.scalar('float', constant_op.constant(2.5))
     mock_scalar_v2.assert_not_called()
 
@@ -71,8 +71,8 @@ def test_scalar_summary_v2__family(self):
     """Tests `family` arg handling when scalar v2 is invoked."""
     with test.mock.patch.object(
         summary_v2, 'scalar', autospec=True) as mock_scalar_v2:
-      with summary_ops_v2.create_summary_file_writer('/tmp/test').as_default(
-          step=1):
+      with summary_ops_v2.create_summary_file_writer(
+          self.get_temp_dir()).as_default(step=1):
         tensor = summary_lib.scalar(
             'float', constant_op.constant(2.5), family='otter')
     # Returns empty string.
@@ -86,8 +86,8 @@ def test_scalar_summary_v2__family_w_outer_scope(self):
     """Tests `family` arg handling when there is an outer scope."""
     with test.mock.patch.object(
         summary_v2, 'scalar', autospec=True) as mock_scalar_v2:
-      with summary_ops_v2.create_summary_file_writer('/tmp/test').as_default(
-          step=1):
+      with summary_ops_v2.create_summary_file_writer(
+          self.get_temp_dir()).as_default(step=1):
         with ops.name_scope_v2('sea'):
           tensor = summary_lib.scalar(
               'float', constant_op.constant(3.5), family='crabnet')
@@ -104,7 +104,8 @@ def test_scalar_summary_v2__v1_set_step(self):
     global_step.assign(1024)
     with test.mock.patch.object(
         summary_v2, 'scalar', autospec=True) as mock_scalar_v2:
-      with summary_ops_v2.create_summary_file_writer('/tmp/test').as_default():
+      with summary_ops_v2.create_summary_file_writer(
+          self.get_temp_dir()).as_default():
         i = constant_op.constant(2.5)
         tensor = summary_lib.scalar('float', i)
     # Returns empty string.
@@ -117,8 +118,8 @@ def test_image_summary_v2(self):
     """Tests image v2 invocation."""
     with test.mock.patch.object(
         summary_v2, 'image', autospec=True) as mock_image_v2:
-      with summary_ops_v2.create_summary_file_writer('/tmp/test').as_default(
-          step=2):
+      with summary_ops_v2.create_summary_file_writer(
+          self.get_temp_dir()).as_default(step=2):
         i = array_ops.ones((5, 4, 4, 3))
         with ops.name_scope_v2('outer'):
           tensor = summary_lib.image('image', i, max_outputs=3, family='family')
@@ -128,6 +129,21 @@ def test_image_summary_v2(self):
     mock_image_v2.assert_called_once_with(
         'family/outer/family/image', data=i, step=2, max_outputs=3)
 
+  @test_util.run_v2_only
+  def test_histogram_summary_v2(self):
+    """Tests histogram v2 invocation."""
+    with test.mock.patch.object(
+        summary_v2, 'histogram', autospec=True) as mock_histogram_v2:
+      with summary_ops_v2.create_summary_file_writer(
+          self.get_temp_dir()).as_default(step=3):
+        i = array_ops.ones((1024,))
+        tensor = summary_lib.histogram('histogram', i, family='family')
+    # Returns empty string.
+    self.assertEqual(tensor.numpy(), b'')
+    self.assertEqual(tensor.dtype, dtypes.string)
+    mock_histogram_v2.assert_called_once_with(
+        'family/family/histogram', data=i, step=3)
+
 
 if __name__ == '__main__':
   test.main()