code: def real_svg_loss(
      top_out,
      targets,
      model_hparams,
      unused_vocab_size,
      unused_weights_fn,
  ):
      """Computes loss for svg decoder model."""
      # targets already come in 10-dim mode, no need to so any mdn stuff
      # obviously.
      targets_commands_rel = targets[..., :4]
      targets_args_rel = targets[..., 4:]
  
      with tf.variable_scope("full_command_loss"):
          num_mix = model_hparams.num_mixture
          commands = top_out[:, :, :, :4]
          args = top_out[:, :, :, 4:]
          # args are [batch, seq, 1, 6*3*num_mix]. want [batch * seq * 6, 3*num_mix]
          args = tf.reshape(args, [-1, 3 * num_mix])
          out_logmix, out_mean, out_logstd = _get_mdn_coef(args)
  
          # before we compute mdn_args_loss, we need to create a mask for elements
          # to ignore on it.
          # create mask
          masktemplate = tf.constant(
              [
                  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                  [0.0, 0.0, 0.0, 0.0, 1.0, 1.0],
                  [0.0, 0.0, 0.0, 0.0, 1.0, 1.0],
                  [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
              ]
          )
          mask = tf.tensordot(
              targets_commands_rel, masktemplate, [[-1], [-2]]
          )
  
          # calculate mdn loss, which auto masks it out
          targs_flat = tf.reshape(targets_args_rel, [-1, 1])
          mdn_loss = _get_mdn_loss(
              out_logmix,
              out_mean,
              out_logstd,
              targs_flat,
              mask,
              model_hparams.dont_reduce_loss,
          )
  
          # we dont have to manually mask out the softmax xent loss because
          # internally, each dimention of the xent loss is multiplied by the
          # given probability in the label for that dim. So for a one-hot label [0,
          # 1, 0] the xent loss between logit[0] and label[0] are multiplied by 0,
          # whereas between logit[1] and label[1] are multiplied by 1. Because our
          # targets_commands_rel is all 0s for the padding, sofmax_xent_loss is 0
          # for those elements as well.
          softmax_xent_loss = tf.nn.softmax_cross_entropy_with_logits(
              labels=targets_commands_rel, logits=commands
          )
  
          # Accumulate losses
          if model_hparams.dont_reduce_loss:
              softmax_xent_loss = tf.reduce_mean(
                  softmax_xent_loss, [1, 2]
              )
          else:
              softmax_xent_loss = tf.reduce_mean(softmax_xent_loss)
          loss = (
              model_hparams.mdn_k * mdn_loss
              + model_hparams.soft_k * softmax_xent_loss
          )
  
      global _summarized_losses
      if not _summarized_losses:
          with tf.name_scope(None), tf.name_scope("losses_command"):
              tf.summary.scalar("mdn_loss", mdn_loss)
              tf.summary.scalar("softmax_xent_loss", softmax_xent_loss)
  
      # this tells us not to re-create the summary ops
      _summarized_losses = True
  
      return loss, tf.constant(1.0)
  
file_path: magenta/magenta/models/svg_vae/svg_decoder_loss.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import modalities
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 3 for calling a global function with example usage: real_svg_loss(top_out, targets, model_hparams, unused_vocab_size, unused_weights_fn) and returns: loss,  tf