def get_predictions(logit_per_label: Array) -> Array:
  """Prediction according to https://www.aclweb.org/anthology/P18-1009.pdf."""
  num_labels = logit_per_label.shape[1]
  # Independent (per-label) predictions
  predictions = (logit_per_label > 0).astype(jnp.int32)
  # A single most confident prediction
  single_best_prediction = jnp.argmax(logit_per_label, axis=-1)
  single_best_prediction = jax.nn.one_hot(
      single_best_prediction, num_labels, dtype=jnp.int32)
  predictions_exists = predictions.sum(axis=1, keepdims=True) > 0
  predictions_exists = predictions_exists.astype(jnp.int32)
  final_predictions = (
      predictions_exists * predictions +
      (1 - predictions_exists) * single_best_prediction)
  return final_predictions
Prediction according to https://www.aclweb.org/anthology/P18-1009.pdf.
language/mentionmemory/tasks/ultra_fine_entity_typing_task.py