code: def test_custom_summary(self):
      def scalar_v2_mock(name, data, step=None):
          """A reimplementation of the scalar plugin to avoid circular
          deps."""
          metadata = SummaryMetadata()
          # Should match value in tensorboard/plugins/scalar/metadata.py.
          metadata.plugin_data.plugin_name = "scalars"
          with summary.experimental.summary_scope(
              name, "scalar_summary", values=[data, step]
          ) as (tag, _):
              tensor = backend.convert_to_tensor(data, dtype="float32")
              if backend.backend() == "torch":
                  # TODO: Use device scope after the API is added.
                  if tensor.is_cuda:
                      tensor = tensor.cpu()
              summary.write(
                  tag=tag,
                  tensor=tensor,
                  step=step,
                  metadata=metadata,
              )
  
      class LayerWithSummary(layers.Layer):
          def call(self, x):
              scalar_v2_mock("custom_summary", ops.sum(x))
              return x
  
      model = models.Sequential(
          [
              layers.Input((5,)),
              LayerWithSummary(),
          ]
      )
  
      # summary ops not compatible with XLA
      model.compile("sgd", "mse", jit_compile=False)
      logdir, train_dir, validation_dir = self._get_log_dirs()
      tb_cbk = callbacks.TensorBoard(logdir, update_freq=1)
      x, y = np.ones((10, 5)), np.ones((10, 5))
      model.fit(
          x, y, batch_size=2, validation_data=(x, y), callbacks=[tb_cbk]
      )
      summary_file = list_summaries(logdir)
      # TODO: tensorflow will tag with model/layer_with_summary/custom_summary
      # Jax will only use custom_summary tag
      self.assertEqual(
          self._strip_to_only_final_name(summary_file.scalars),
          {
              _ObservedSummary(logdir=train_dir, tag="batch_loss"),
              _ObservedSummary(logdir=train_dir, tag="epoch_loss"),
              _ObservedSummary(
                  logdir=train_dir, tag="epoch_learning_rate"
              ),
              _ObservedSummary(logdir=validation_dir, tag="epoch_loss"),
              _ObservedSummary(
                  logdir=validation_dir,
                  tag="evaluation_loss_vs_iterations",
              ),
              _ObservedSummary(
                  logdir=train_dir,
                  tag="custom_summary",
              ),
              _ObservedSummary(
                  logdir=validation_dir,
                  tag="custom_summary",
              ),
          },
      )
  
imports: import collections
  import os
  import random
  import numpy
  import pytest
  import tensorflow.summary
  from tensorflow.compat.v1 import SummaryMetadata
  from tensorflow.core.util import event_pb2
  from tensorflow.python.lib.io import tf_record
  from keras import backend
  from keras import callbacks
  from keras import layers
  from keras import losses
  from keras import models
  from keras import ops
  from keras import optimizers
  from keras import testing
  from keras.optimizers import schedules
label_desc: Use image node_id 12 for calling the TestTensorBoardV2 obj's underlying member method code with example usage: obj.test_custom_summary() and returns: x