def _write_keras_model_train_graph(self):
        """Writes Keras model train_function graph to TensorBoard."""
        with self._train_writer.as_default():
            train_fn = self.model.train_function
            # If the train_function is a `tf.function`, we can write out a
            # graph
            if hasattr(train_fn, "function_spec"):
                # TODO(b/243822285): Use _variable_creation_fn directly.
                if hasattr(train_fn, "_concrete_stateful_fn"):
                    self.summary.graph(train_fn._concrete_stateful_fn.graph)
                else:
                    self.summary.graph(
                        train_fn._concrete_variable_creation_fn.graph
                    )
Writes Keras model train_function graph to TensorBoard.
keras/src/callbacks/tensorboard.py