def WeightedFScore(beta=1., initial_category_index=0):
  """Returns a layer that computes a weighted F-score.

  The weighted F-score summarize how well the classifier's `k` predictions
  align with the observed/gold instances of `k`. It additionally
  weights the summary by the number of observed/gold and predicted examples
  in each class.

  Args:
    beta: a parameter that determines the weight of recall in the F-score.
    initial_category_index: an index of the initial category.

  The layer takes two inputs:

    - Model output from one batch, an ndarray of float-valued elements.

    - A batch of element-wise target values, which matches the shape of the
      model output.

  The layer returns a weighted F-score across all the classes.
  """
  def f(model_output, targets):  # pylint: disable=invalid-name
    beta2 = beta ** 2
    predictions = jnp.argmax(model_output, axis=-1)
    n_categories = model_output.shape[-1]
    f_scores = jnp.empty(0)
    weights = jnp.empty(0)
    for k in range(initial_category_index, n_categories):
      _, _, n_k_targets, precision, recall = _precision_recall(
          predictions, targets, k)
      f_scores = jnp.append(f_scores, _f_score(precision, recall, beta2))
      weights = jnp.append(weights, n_k_targets)
    return jnp.average(f_scores, weights=weights)

  return base.Fn('WeightedFScore', f)
Returns a layer that computes a weighted F-score.

  The weighted F-score summarize how well the classifier's `k` predictions
  align with the observed/gold instances of `k`. It additionally
  weights the summary by the number of observed/gold and predicted examples
  in each class.

  Args:
    beta: a parameter that determines the weight of recall in the F-score.
    initial_category_index: an index of the initial category.

  The layer takes two inputs:

    - Model output from one batch, an ndarray of float-valued elements.

    - A batch of element-wise target values, which matches the shape of the
      model output.

  The layer returns a weighted F-score across all the classes.
  
trax/layers/metrics.py