code: def smoothing_cross_entropy_factored_grad(op, dy):
    
    a = op.inputs[0]
    b = op.inputs[1]
    labels = op.inputs[2]
    confidence = op.inputs[3]
    num_splits = 16
    vocab_size = shape_list(b)[0]
    labels = approximate_split(labels, num_splits)
    a = approximate_split(a, num_splits)
    dy = approximate_split(dy, num_splits)
    b_grad = None
    a_grad_parts = []
    deps = []
    for part in range(num_splits):
      with tf.control_dependencies(deps):
        logits = tf.matmul(a[part], b, transpose_b=True)
        output_part = smoothing_cross_entropy(logits, labels[part], vocab_size,
                                              confidence)
        a_grad_part, b_grad_part = tf.gradients(
            ys=[output_part], xs=[a[part], b], grad_ys=[dy[part]])
        a_grad_parts.append(a_grad_part)
        if part > 0:
          b_grad += b_grad_part
        else:
          b_grad = b_grad_part
        deps = [b_grad, a_grad_part]
    a_grad = tf.concat(a_grad_parts, 0)
    return a_grad, b_grad, None, None
docs: Gradient function for smoothing_cross_entropy_factored.