code: def smoothing_cross_entropy_factored(a, b, labels, confidence):
    
    num_splits = 16
    vocab_size = shape_list(b)[0]
    labels = approximate_split(labels, num_splits)
    a = approximate_split(a, num_splits)
    parts = []
    for part in range(num_splits):
      with tf.control_dependencies(parts[-1:]):
        logits = tf.matmul(a[part], b, transpose_b=True)
        parts.append(
            smoothing_cross_entropy(logits, labels[part], vocab_size, confidence))
    return tf.concat(parts, 0)
docs: Memory-efficient computation of smoothing cross-entropy.
  
    Avoids realizing the entire logits matrix at once.
  
    Args:
      a: a Tensor with shape [batch, inner_dim]
      b: a Tensor with shape [vocab_size, inner_dim]
      labels: an integer Tensor with shape [batch]
      confidence: a float
  
    Returns:
      A Tensor with shape [batch]