def prefix_accuracy(predictions,
                    labels,
                    weights_fn=common_layers.weights_nonzero):
  """Average # of correct tokens at start of sequences, ignoring padding 0s.

  See section 4.3 of Learning to Transduce with Unbounded Memory,
  Grefenstette et al., 2015.

  Args:
    predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
        type tf.float32 representing the logits, 0-padded.
    labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
        representing the labels of same length as logits and 0-padded.
    weights_fn: ignored. The weights returned are the total length of the ground
        truth labels, excluding 0-paddings.

  Returns:
    (prefix accuracy, 1.0)

  Raises:
    ValueError: if weights_fn is not common_layers.weights_nonzero.
  """
  if weights_fn is not common_layers.weights_nonzero:
    raise ValueError("Only weights_nonzero can be used for this metric.")

  predictions = tf.to_int32(tf.squeeze(tf.argmax(predictions, axis=-1), axis=2))
  labels = tf.squeeze(labels, axis=(2, 3))
  seq_len = tf.reduce_sum(
      tf.cast(tf.not_equal(labels, tf.constant(0)), dtype=tf.float32), axis=1)
  matching_elements = tf.equal(labels, predictions)
  prefix_len = tf.reduce_sum(
      tf.cumprod(tf.cast(matching_elements, tf.float32), axis=1), axis=1)
  return tf.reduce_mean(prefix_len / seq_len), tf.constant(1.0)
Average # of correct tokens at start of sequences, ignoring padding 0s.

  See section 4.3 of Learning to Transduce with Unbounded Memory,
  Grefenstette et al., 2015.

  Args:
    predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
        type tf.float32 representing the logits, 0-padded.
    labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
        representing the labels of same length as logits and 0-padded.
    weights_fn: ignored. The weights returned are the total length of the ground
        truth labels, excluding 0-paddings.

  Returns:
    (prefix accuracy, 1.0)

  Raises:
    ValueError: if weights_fn is not common_layers.weights_nonzero.
  
tensor2tensor/utils/metrics.py