code: def padded_cross_entropy_factored(factored_logits,
                                    labels,
                                    label_smoothing,
                                    weights_fn=weights_nonzero,
                                    reduce_sum=True):
    
    a = factored_logits.a
    b = factored_logits.b
    confidence = 1.0 - label_smoothing
    with tf.name_scope("padded_cross_entropy_factored", values=[a, b, labels]):
      labels_flat = tf.reshape(labels, [-1])
      a_flat = tf.reshape(a, [-1, shape_list(b)[1]])
      xent = smoothing_cross_entropy_factored(a_flat, b, labels_flat,
                                              tf.convert_to_tensor(confidence))
      xent = tf.reshape(xent, shape_list(labels))
      weights = weights_fn(labels)
      if not reduce_sum:
        return xent * weights, weights
      return tf.reduce_sum(xent * weights), tf.reduce_sum(weights)
docs: Memory-efficient computation of smoothing cross-entropy.
  
    Avoids realizing the entire logits matrix at once.
  
    Args:
      factored_logits: a `FactoredTensor` representing a Tensor
         with shape `[batch, timesteps, vocab_size]`.
      labels: an integer `Tensor` with shape `[batch, timesteps]`.
      label_smoothing: a floating point `Scalar`.
      weights_fn: A function from labels to weights.
      reduce_sum: a Boolean, whether to sum at the end or not.
  
    Returns:
      loss_numerator: a `Scalar`.  Sum of losses.
      loss_denominator: a `Scalar.  The number of non-padding target tokens.