note this file and tensorboard in general has a dependency on tensorflow a summary that was emitted during a test fields logdir str the logdir of the filewriter to which the summary was written tag str the name of the summary yields event protocol buffers from a given path def initself path self tfrecorditerator tfrecord tfrecorditeratorpath def iterself return self def nextself r nextself tfrecorditerator return eventpb2 event fromstringr next next class summaryfile def initself self scalars set self images set self histograms set self tensors set self graphdefs self convertfromv2summaryproto false def listsummarieslogdir result summaryfile for dirpath filenames in os walklogdir for filename in filenames if not filename startswithevents out continue path os path joindirpath filename for event in summaryiteratorpath if event graphdef result graphdefs appendevent graphdef if not event summary e g it s a graphdef event continue for value in event summary value tag value tag case on the value rather than the summary metadata because the keras callback uses summaryopsv2 to emit oldstyle summaries see b124535134 kind value whichoneofvalue container simplevalue result scalars image result images histo result histograms tensor result tensors getkind if container is none raise valueerror unexpected summary kind r in event file s nr kind path event elif kind tensor and tag keras convert the tf2 summary proto to old style for type checking pluginname value metadata plugindata pluginname container images result images histograms result histograms scalars result scalars getpluginname if container is not none result convertfromv2summaryproto true else container result tensors container addobservedsummarylogdirdirpath tagtag return result class testtensorboardv2testing testcase def getlogdirsself logdir os path join self gettempdir strrandom randint1 1e7 tb traindir os path joinlogdir train validationdir os path joinlogdir validation return logdir traindir validationdir def getmodelself compilemodeltrue model models sequential layers input10 10 1 layers flatten layers dense1 if compilemodel model compilesgd mse return model pytest mark requirestrainablebackend def testtensorboardbasicself model self getmodel x y np ones10 10 10 1 np ones10 1 logdir traindir validationdir self getlogdirs tbcbk callbacks tensorboardlogdir model fit x y batchsize2 epochs2 validationdatax y callbackstbcbk summaryfile listsummarieslogdir self assertequal summaryfile scalars observedsummarylogdirtraindir tagepochloss observedsummarylogdirvalidationdir tagepochloss observedsummarylogdirtraindir tagepochlearningrate observedsummary logdirvalidationdir tagevaluationlossvsiterations pytest mark requirestrainablebackend def testtensorboardacrossinvocationsself a reimplementation of the scalar plugin to avoid circular deps metadata summarymetadata should match value in tensorboardpluginsscalarmetadata py metadata plugindata pluginname scalars with summary experimental summaryscope name scalarsummary valuesdata step as tag tensor backend converttotensordata dtypefloat32 if backend backend torch todo use device scope after the api is added if tensor iscuda tensor tensor cpu summary write tagtag tensortensor stepstep metadatametadata class layerwithsummarylayers layer def callself x scalarv2mockcustomsummary ops sumx return x model models sequential layers input5 layerwithsummary summary ops not compatible with xla model compilesgd mse jitcompilefalse logdir traindir validationdir self getlogdirs tbcbk callbacks tensorboardlogdir updatefreq1 x y np ones10 5 np ones10 5 model fit x y batchsize2 validationdatax y callbackstbcbk summaryfile listsummarieslogdir todo tensorflow will tag with modellayerwithsummarycustomsummary jax will only use customsummary tag self assertequal self striptoonlyfinalnamesummaryfile scalars observedsummarylogdirtraindir tagbatchloss observedsummarylogdirtraindir tagepochloss observedsummarylogdirtraindir tagepochlearningrate observedsummarylogdirvalidationdir tagepochloss observedsummary logdirvalidationdir tagevaluationlossvsiterations observedsummary logdirtraindir tagcustomsummary observedsummary logdirvalidationdir tagcustomsummary self assertequal summaryfile scalars observedsummarylogdirtraindir tagbatchloss observedsummarylogdirtraindir tagepochloss observedsummarylogdirvalidationdir tagepochloss observedsummary logdirvalidationdir tagevaluationlossvsiterations observedsummary logdirtraindir tagmodellayerwithsummarycustomsummary observedsummary logdirvalidationdir tagmodellayerwithsummarycustomsummary def striptoonlyfinalnameself summaries result set for s in summaries if not in s tag result adds else newtag s tag split1 result adds replacetagnewtag return result def striplayernamesself summaries modeltype result set for s in summaries if not in s tag raise valueerrorftag has no layer name s tag r startfrom 2 if subclass in modeltype else 1 newtag joins tag splitstartfrom result adds replacetagnewtag return result def stripvariablenamesself summaries result set for s in summaries if not in s tag result adds else splittag s tag split if variable in splittag0 result adds replacetagsplittag1 else result adds return result pytest mark skipif backend backend torch reasontorch backend requires blocking numpy conversion pytest mark requirestrainablebackend def testtensorboardnonblockingself logdir self getlogdirs model models sequentiallayers dense1 model optimizer optimizers adam tb callbacks tensorboardlogdir cblist callbacks callbacklist tb modelmodel epochs1 steps100 verbose0 tensor ops converttotensor1 0 def mocknumpy raise runtimeerror if this error is seen tensorboard is causing a blocking numpy conversion tensor numpy mocknumpy logs metric tensor cblist ontrainbeginlogs cblist onepochbegin0 logs cblist ontrainbatchbegin0 logs cblist ontrainbatchend0 logs cblist onepochend0 logs cblist ontrainendlogs cblist ontestbeginlogs cblist ontestbatchbegin0 logs cblist ontestbatchend0 logs cblist ontestendlogs cblist onpredictbeginlogs cblist onpredictbatchbeginlogs cblist onpredictbatchendlogs cblist onpredictendlogs def countxplanefileself logdir profiledir os path joinlogdir plugins profile count 0 for dirpath dirnames filenames in os walkprofiledir del dirpath unused del dirnames unused for filename in filenames if filename endswith xplane pb count 1 return count def fitmodelandassertkerasmodelwrittenself model x y np ones10 10 10 1 np ones10 1 logdir traindir validationdir self getlogdirs tbcbk callbacks tensorboard logdir writegraphtrue profilebatch0 model fit x y batchsize2 epochs3 validationdatax y callbackstbcbk summaryfile listsummarieslogdir self assertequal summaryfile tensors observedsummarylogdirtraindir tagkeras if not model runeagerly there should be one train graph self assertlensummaryfile graphdefs 1 for graphdef in summaryfile graphdefs graphdefstr strgraphdef all the model layers should appear in the graphs for layer in model layers if input not in layer name self assertinlayer name graphdefstr def testtensorboardwritesequentialmodelnoinputshapeself todo requires tojson implementation in trainer model models sequential conv2d8 3 3 flatten dense1 model compilesgd mse self fitmodelandassertkerasmodelwrittenmodel pass def testtensorboardwritesequentialmodelwithinputshapeself todo requires tojson implementation in trainer model models sequential inputinputshape10 10 1 conv2d8 3 3 flatten dense1 model compilesgd mse self fitmodelandassertkerasmodelwrittenmodel pass def testtensorboardwritemodelself todo requires tojson implementation in trainer see https github comkerasteamkerasblob a8d4a7f1ffc9de3c5932828a107e4e95e8803fb4 kerasenginetraining pyl3313 inputs input10 10 1 x conv2d8 3 3 activationreluinputs x flattenx x dense1x model models modelinputsinputs outputsx model compilesgd mse breakpoint self fitmodelandassertkerasmodelwrittenmodel pass pytest mark skipif backend backend tensorflow reasonthe profiling test can only run with tf backend def testtensorboardautotraceself todo waiting for implementation for torchjax for profiling ops if backend backend jax return todo debug profiling for jax logdir traindir validationdir self getlogdirs model models sequential layers input10 10 1 layers flatten layers dense1 x y np ones10 10 10 1 np ones10 1 tbcbk callbacks tensorboard logdir histogramfreq1 profilebatch1 writegraphfalse model compilesgd mse model fit x y batchsize2 epochs2 validationdatax y callbackstbcbk summaryfile listsummarieslogdir self assertequal summaryfile tensors observedsummarylogdirtraindir tagbatch1 self assertequal1 self countxplanefilelogdirlogdir pass note this file and tensorboard in general has a dependency on tensorflow a summary that was emitted during a test fields logdir str the logdir of the filewriter to which the summary was written tag str the name of the summary yields event protocol buffers from a given path a record of summary tags and the files to which they were written fields scalars images histograms and tensors are sets containing _observedsummary values read all summaries under the logdir into a _summaryfile args logdir a path to a directory that contains zero or more event files either as direct children or in transitive subdirectories summaries in these events must only contain old style scalars images and histograms non summary events like graph_def s are ignored returns a _summaryfile object reflecting all summaries written to any event files in the logdir or any of its descendant directories raises valueerror if an event file contains an summary of unexpected kind e g it s a graph_def event case on the value rather than the summary metadata because the keras callback uses summary_ops_v2 to emit old style summaries see b 124535134 convert the tf2 summary proto to old style for type checking regression test for summary writer resource use after free a reimplementation of the scalar plugin to avoid circular deps should match value in tensorboard plugins scalar metadata py todo use device scope after the api is added summary ops not compatible with xla todo tensorflow will tag with model layer_with_summary custom_summary jax will only use custom_summary tag self assertequal summary_file scalars _observedsummary logdir train_dir tag batch_loss _observedsummary logdir train_dir tag epoch_loss _observedsummary logdir validation_dir tag epoch_loss _observedsummary logdir validation_dir tag evaluation_loss_vs_iterations _observedsummary logdir train_dir tag model layer_with_summary custom_summary _observedsummary logdir validation_dir tag model layer_with_summary custom_summary removes all leading names in a summary args summaries a set of _observedsummary values returns a new set of _observedsummary values striped of all name except for the terminal one deduplicate summary names modulo layer prefix this removes the first slash component of each tag name for instance foo bar baz becomes bar baz args summaries a set of _observedsummary values model_type the model type currently being tested returns a new set of _observedsummary values with layer prefixes removed remove variable_n from summary tag variable_n tag names are added with random numbers removing them ensures deterministic tag names args summaries a set of _observedsummary values returns a new set of _observedsummary values with layer prefixes removed unused unused there should be one train graph all the model layers should appear in the graphs todo requires to_json implementation in trainer model models sequential conv2d 8 3 3 flatten dense 1 model compile sgd mse self fitmodelandassertkerasmodelwritten model todo requires to_json implementation in trainer model models sequential input input_shape 10 10 1 conv2d 8 3 3 flatten dense 1 model compile sgd mse self fitmodelandassertkerasmodelwritten model todo requires to_json implementation in trainer see https github com keras team keras blob a8d4a7f1ffc9de3c5932828a107e4e95e8803fb4 keras engine training py l3313 inputs input 10 10 1 x conv2d 8 3 3 activation relu inputs x flatten x x dense 1 x model models model inputs inputs outputs x model compile sgd mse breakpoint self fitmodelandassertkerasmodelwritten model todo waiting for implementation for torch jax for profiling ops if backend backend jax return todo debug profiling for jax
import collections
import os
import random

import numpy as np
import pytest
import tensorflow.summary as 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







_ObservedSummary = collections.namedtuple("_ObservedSummary", ("logdir", "tag"))


class _SummaryIterator:
    

    def __init__(self, path):
        self._tf_record_iterator = tf_record.tf_record_iterator(path)

    def __iter__(self):
        return self

    def __next__(self):
        r = next(self._tf_record_iterator)
        return event_pb2.Event.FromString(r)

    next = __next__


class _SummaryFile:
    

    def __init__(self):
        self.scalars = set()
        self.images = set()
        self.histograms = set()
        self.tensors = set()
        self.graph_defs = []
        self.convert_from_v2_summary_proto = False


def list_summaries(logdir):
    
    result = _SummaryFile()
    for dirpath, _, filenames in os.walk(logdir):
        for filename in filenames:
            if not filename.startswith("events.out."):
                continue
            path = os.path.join(dirpath, filename)
            for event in _SummaryIterator(path):
                if event.graph_def:
                    result.graph_defs.append(event.graph_def)
                if not event.summary:  
                    continue
                for value in event.summary.value:
                    tag = value.tag
                    
                    
                    
                    kind = value.WhichOneof("value")
                    container = {
                        "simple_value": result.scalars,
                        "image": result.images,
                        "histo": result.histograms,
                        "tensor": result.tensors,
                    }.get(kind)
                    if container is None:
                        raise ValueError(
                            "Unexpected summary kind %r in event file %s:\n%r"
                            % (kind, path, event)
                        )
                    elif kind == "tensor" and tag != "keras":
                        
                        
                        plugin_name = value.metadata.plugin_data.plugin_name
                        container = {
                            "images": result.images,
                            "histograms": result.histograms,
                            "scalars": result.scalars,
                        }.get(plugin_name)
                        if container is not None:
                            result.convert_from_v2_summary_proto = True
                        else:
                            container = result.tensors
                    container.add(_ObservedSummary(logdir=dirpath, tag=tag))
    return result


class TestTensorBoardV2(testing.TestCase):
    def _get_log_dirs(self):
        logdir = os.path.join(
            self.get_temp_dir(), str(random.randint(1, 1e7)), "tb"
        )
        train_dir = os.path.join(logdir, "train")
        validation_dir = os.path.join(logdir, "validation")
        return logdir, train_dir, validation_dir

    def _get_model(self, compile_model=True):
        model = models.Sequential(
            [
                layers.Input((10, 10, 1)),
                layers.Flatten(),
                layers.Dense(1),
            ]
        )
        if compile_model:
            model.compile("sgd", "mse")
        return model

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_basic(self):
        model = self._get_model()
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))
        logdir, train_dir, validation_dir = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(logdir)

        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            validation_data=(x, y),
            callbacks=[tb_cbk],
        )

        summary_file = list_summaries(logdir)
        self.assertEqual(
            summary_file.scalars,
            {
                _ObservedSummary(logdir=train_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=validation_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=train_dir, tag="epoch_learning_rate"),
                _ObservedSummary(
                    logdir=validation_dir,
                    tag="evaluation_loss_vs_iterations",
                ),
            },
        )

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_across_invocations(self):
        
        model = self._get_model()
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))
        logdir, train_dir, validation_dir = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(logdir)

        for _ in (1, 2):
            model.fit(
                x,
                y,
                batch_size=2,
                epochs=2,
                validation_data=(x, y),
                callbacks=[tb_cbk],
            )

        summary_file = list_summaries(logdir)
        self.assertEqual(
            summary_file.scalars,
            {
                _ObservedSummary(logdir=train_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=validation_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=train_dir, tag="epoch_learning_rate"),
                _ObservedSummary(
                    logdir=validation_dir,
                    tag="evaluation_loss_vs_iterations",
                ),
            },
        )

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_no_spurious_event_files(self):
        model = self._get_model()
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))
        logdir, train_dir, _ = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(logdir)
        model.fit(x, y, batch_size=2, epochs=2, callbacks=[tb_cbk])

        events_file_run_basenames = set()
        for dirpath, _, filenames in os.walk(train_dir):
            if any(fn.startswith("events.out.") for fn in filenames):
                events_file_run_basenames.add(os.path.basename(dirpath))
        self.assertEqual(events_file_run_basenames, {"train"})

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_batch_metrics(self):
        model = self._get_model()
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))
        logdir, train_dir, validation_dir = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(logdir, update_freq=1)

        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            validation_data=(x, y),
            callbacks=[tb_cbk],
        )

        summary_file = list_summaries(logdir)
        self.assertEqual(
            summary_file.scalars,
            {
                _ObservedSummary(logdir=train_dir, tag="batch_loss"),
                _ObservedSummary(logdir=train_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=validation_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=train_dir, tag="epoch_learning_rate"),
                _ObservedSummary(
                    logdir=validation_dir,
                    tag="evaluation_loss_vs_iterations",
                ),
            },
        )

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_learning_rate_schedules(self):
        model = self._get_model(compile_model=False)
        opt = optimizers.SGD(schedules.CosineDecay(0.01, 1))
        model.compile(opt, "mse")
        logdir, train_dir, _ = self._get_log_dirs()
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))

        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            callbacks=[callbacks.TensorBoard(logdir)],
        )

        summary_file = list_summaries(logdir)
        self.assertEqual(
            summary_file.scalars,
            {
                _ObservedSummary(logdir=train_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=train_dir, tag="epoch_learning_rate"),
            },
        )

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_global_step(self):
        model = self._get_model(compile_model=False)
        opt = optimizers.SGD(schedules.CosineDecay(0.01, 1))
        model.compile(opt, "mse")
        logdir, train_dir, _ = self._get_log_dirs()
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))

        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            verbose=0,
            callbacks=[
                callbacks.TensorBoard(
                    logdir,
                    update_freq=1,
                    profile_batch=0,
                    write_steps_per_second=True,
                )
            ],
        )

        summary_file = list_summaries(logdir)
        self.assertEqual(
            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=train_dir, tag="epoch_steps_per_second"
                ),
                _ObservedSummary(
                    logdir=train_dir, tag="batch_steps_per_second"
                ),
            },
        )

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_weight_histograms(self):
        model = self._get_model()
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))
        logdir, train_dir, validation_dir = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(logdir, histogram_freq=1)

        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            validation_data=(x, y),
            callbacks=[tb_cbk],
        )
        summary_file = list_summaries(logdir)

        self.assertEqual(
            summary_file.scalars,
            {
                _ObservedSummary(logdir=train_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=validation_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=train_dir, tag="epoch_learning_rate"),
                _ObservedSummary(
                    logdir=validation_dir,
                    tag="evaluation_loss_vs_iterations",
                ),
            },
        )
        self.assertEqual(
            self._strip_layer_names(summary_file.histograms, "sequential"),
            {_ObservedSummary(logdir=train_dir, tag="histogram")},
        )

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_weight_images(self):
        if backend.config.image_data_format() == "channels_last":
            input_shape = (10, 10, 1)
            x_shape = (10, 10, 10, 1)
        else:
            input_shape = (1, 10, 10)
            x_shape = (10, 1, 10, 10)
        x, y = np.ones(x_shape), np.ones((10, 1))
        logdir, train_dir, validation_dir = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(
            logdir, histogram_freq=1, write_images=True
        )
        model_type = "sequential"
        model = models.Sequential(
            [
                layers.Input(input_shape),
                layers.Conv2D(3, 10),
                layers.GlobalAveragePooling2D(),
                layers.Dense(1),
            ]
        )
        model.compile("sgd", "mse")
        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            validation_data=(x, y),
            callbacks=[tb_cbk],
        )
        summary_file = list_summaries(logdir)

        self.assertEqual(
            summary_file.scalars,
            {
                _ObservedSummary(logdir=train_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=validation_dir, tag="epoch_loss"),
                _ObservedSummary(logdir=train_dir, tag="epoch_learning_rate"),
                _ObservedSummary(
                    logdir=validation_dir,
                    tag="evaluation_loss_vs_iterations",
                ),
            },
        )
        self.assertEqual(
            self._strip_layer_names(summary_file.histograms, model_type),
            {
                _ObservedSummary(logdir=train_dir, tag="histogram"),
            },
        )
        expected_image_summaries = {
            _ObservedSummary(logdir=train_dir, tag="bias/image"),
            _ObservedSummary(logdir=train_dir, tag="kernel/image"),
        }
        self.assertEqual(
            self._strip_variable_names(summary_file.images),
            expected_image_summaries,
        )

    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_projector_callback(self):
        model = models.Sequential(
            [
                layers.Input((10,)),
                layers.Embedding(10, 10, name="test_embedding"),
                layers.Dense(1, activation="sigmoid"),
            ]
        )
        model.compile(
            optimizer="adam", loss=losses.BinaryCrossentropy(from_logits=True)
        )
        x, y = np.ones((10, 10)), np.ones((10, 10))
        logdir, _, _ = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(
            logdir,
            embeddings_freq=1,
            embeddings_metadata={"test_embedding": "metadata.tsv"},
        )

        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            validation_data=(x, y),
            callbacks=[tb_cbk],
        )

        with open(os.path.join(logdir, "projector_config.pbtxt")) as f:
            self.assertEqual(
                f.readlines(),
                [
                    "embeddings {\n",
                    "  tensor_name: "
                    '"layer_with_weights-0/embeddings/.ATTRIBUTES/'
                    'VARIABLE_VALUE"\n',
                    '  metadata_path: "metadata.tsv"\n',
                    "}\n",
                ],
            )

    @pytest.mark.requires_trainable_backend
    def test_custom_summary(self):
        def scalar_v2_mock(name, data, step=None):
            
            metadata = SummaryMetadata()
            
            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":
                    
                    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(),
            ]
        )

        
        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)
        
        
        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",
                ),
            },
        )
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        

    def _strip_to_only_final_name(self, summaries):
        
        result = set()
        for s in summaries:
            if "/" not in s.tag:
                result.add(s)
            else:
                new_tag = s.tag.split("/")[-1]
                result.add(s._replace(tag=new_tag))
        return result

    def _strip_layer_names(self, summaries, model_type):
        
        result = set()
        for s in summaries:
            if "/" not in s.tag:
                raise ValueError(f"tag has no layer name: {s.tag!r}")
            start_from = 2 if "subclass" in model_type else 1
            new_tag = "/".join(s.tag.split("/")[start_from:])
            result.add(s._replace(tag=new_tag))
        return result

    def _strip_variable_names(self, summaries):
        
        result = set()
        for s in summaries:
            if "/" not in s.tag:
                result.add(s)
            else:
                split_tag = s.tag.split("/")
                if "variable" in split_tag[0]:
                    result.add(s._replace(tag=split_tag[-1]))
                else:
                    result.add(s)
        return result

    @pytest.mark.skipif(
        backend.backend() == "torch",
        reason="Torch backend requires blocking numpy conversion.",
    )
    @pytest.mark.requires_trainable_backend
    def test_TensorBoard_non_blocking(self):
        logdir, _, _ = self._get_log_dirs()
        model = models.Sequential([layers.Dense(1)])
        model.optimizer = optimizers.Adam()
        tb = callbacks.TensorBoard(logdir)
        cb_list = callbacks.CallbackList(
            [tb], model=model, epochs=1, steps=100, verbose=0
        )
        tensor = ops.convert_to_tensor(1.0)

        def mock_numpy():
            raise RuntimeError(
                "If this error is seen, TensorBoard is causing a blocking "
                "NumPy conversion."
            )

        tensor.numpy = mock_numpy

        logs = {"metric": tensor}

        cb_list.on_train_begin(logs)
        cb_list.on_epoch_begin(0, logs)
        cb_list.on_train_batch_begin(0, logs)
        cb_list.on_train_batch_end(0, logs)
        cb_list.on_epoch_end(0, logs)
        cb_list.on_train_end(logs)

        cb_list.on_test_begin(logs)
        cb_list.on_test_batch_begin(0, logs)
        cb_list.on_test_batch_end(0, logs)
        cb_list.on_test_end(logs)

        cb_list.on_predict_begin(logs)
        cb_list.on_predict_batch_begin(logs)
        cb_list.on_predict_batch_end(logs)
        cb_list.on_predict_end(logs)

    def _count_xplane_file(self, logdir):
        profile_dir = os.path.join(logdir, "plugins", "profile")
        count = 0
        for dirpath, dirnames, filenames in os.walk(profile_dir):
            del dirpath  
            del dirnames  
            for filename in filenames:
                if filename.endswith(".xplane.pb"):
                    count += 1
        return count

    def fitModelAndAssertKerasModelWritten(self, model):
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))
        logdir, train_dir, validation_dir = self._get_log_dirs()
        tb_cbk = callbacks.TensorBoard(
            logdir, write_graph=True, profile_batch=0
        )
        model.fit(
            x,
            y,
            batch_size=2,
            epochs=3,
            validation_data=(x, y),
            callbacks=[tb_cbk],
        )
        summary_file = list_summaries(logdir)
        self.assertEqual(
            summary_file.tensors,
            {
                _ObservedSummary(logdir=train_dir, tag="keras"),
            },
        )
        if not model.run_eagerly:
            
            self.assertLen(summary_file.graph_defs, 1)
            for graph_def in summary_file.graph_defs:
                graph_def_str = str(graph_def)

                
                for layer in model.layers:
                    if "input" not in layer.name:
                        self.assertIn(layer.name, graph_def_str)

    def test_TensorBoard_write_sequential_model_no_input_shape(self):
        
        
        
        
        
        
        
        
        
        
        pass

    def test_TensorBoard_write_sequential_model_with_input_shape(self):
        
        
        
        
        
        
        
        
        
        
        
        pass

    def test_TensorBoard_write_model(self):
        
        
        
        
        
        
        
        
        
        
        
        
        pass

    @pytest.mark.skipif(
        backend.backend() != "tensorflow",
        reason="The profiling test can only run with TF backend.",
    )
    def test_TensorBoard_auto_trace(self):
        
        
        
        
        logdir, train_dir, validation_dir = self._get_log_dirs()
        model = models.Sequential(
            [
                layers.Input((10, 10, 1)),
                layers.Flatten(),
                layers.Dense(1),
            ]
        )
        x, y = np.ones((10, 10, 10, 1)), np.ones((10, 1))
        tb_cbk = callbacks.TensorBoard(
            logdir, histogram_freq=1, profile_batch=1, write_graph=False
        )
        model.compile("sgd", "mse")
        model.fit(
            x,
            y,
            batch_size=2,
            epochs=2,
            validation_data=(x, y),
            callbacks=[tb_cbk],
        )
        summary_file = list_summaries(logdir)

        self.assertEqual(
            summary_file.tensors,
            {
                _ObservedSummary(logdir=train_dir, tag="batch_1"),
            },
        )
        self.assertEqual(1, self._count_xplane_file(logdir=logdir))
        pass