code: def set_precision(predictions, labels,
                    weights_fn=common_layers.weights_nonzero):
    
    with tf.variable_scope("set_precision", values=[predictions, labels]):
      labels = tf.squeeze(labels, [2, 3])
      weights = weights_fn(labels)
      labels = tf.one_hot(labels, predictions.shape[-1])
      labels = tf.reduce_max(labels, axis=1)
      labels = tf.cast(labels, tf.bool)
      return tf.to_float(tf.equal(labels, predictions)), weights
docs: Precision of set predictions.
  
    Args:
      predictions : A Tensor of scores of shape [batch, nlabels].
      labels: A Tensor of int32s giving true set elements,
        of shape [batch, seq_length].
      weights_fn: A function to weight the elements.
  
    Returns:
      hits: A Tensor of shape [batch, nlabels].
      weights: A Tensor of shape [batch, nlabels].