code: def style_loss(style_gram_matrices, end_points, style_weights):
      """Style loss.
  
      Args:
        style_gram_matrices: dict mapping VGG16 layer names to their corresponding
            gram matrix for the style image.
        end_points: dict mapping VGG16 layer names to their corresponding
            Tensor value for the stylized input.
        style_weights: dict mapping layer names to their associated style loss
            weight. Keys that are missing from the dict won't have their style
            loss computed.
  
      Returns:
        Tensor for the total style loss, dict mapping loss names to losses.
      """
      total_style_loss = np.float32(0.0)
      style_loss_dict = {}
  
      for name in style_weights:
          # Reducing over all but the batch axis before multiplying with the style
          # weights allows to use multiple sets of style weights in a single batch.
          loss = tf.reduce_mean(
              (
                  gram_matrix(end_points[name])
                  - style_gram_matrices[name]
              )
              ** 2,
              [1, 2],
          )
          weighted_style_loss = tf.reduce_mean(
              style_weights[name] * loss
          )
          loss = tf.reduce_mean(loss)
  
          style_loss_dict["style_loss/" + name] = loss
          style_loss_dict[
              "weighted_style_loss/" + name
          ] = weighted_style_loss
          total_style_loss += weighted_style_loss
  
      style_loss_dict["total_style_loss"] = total_style_loss
  
      return total_style_loss, style_loss_dict
  
file_path: magenta/magenta/models/image_stylization/learning.py
imports: from magenta.models.image_stylization import vgg
  import numpy
  import tensorflow.compat.v1
  import tf_slim
label_desc: Use image node_id 4 for calling a global function with example usage: style_loss(style_gram_matrices, end_points, style_weights) and returns: total_style_loss,  style_loss_dict