def _combine_multi_example_logits(logits):
    """Combine per-example logits into a per-class logit."""
    logits = tf.reshape(
        logits, [-1, fewshot_num_classes, fewshot_num_examples_per_class, 1])
    if class_examples_combiner == "max":
      logits = tf.reduce_max(logits, axis=2)
    if class_examples_combiner == "mean":
      logits = tf.reduce_mean(logits, axis=2)
    if class_examples_combiner == "logsumexp":
      logits = tf.reduce_logsumexp(logits, axis=2)
    if class_examples_combiner == "min":
      logits = tf.reduce_min(logits, axis=2)
    if class_examples_combiner == "sigmoid_mean":
      logits = tf.sigmoid(logits)
      logits = tf.reduce_mean(logits, axis=2)
    return logits
Combine per-example logits into a per-class logit.
language/relation_learning/models/bert_fewshot_classifier.py