def MacroAveragedFScore(beta=1., initial_category_index=0):
  r"""Returns a layer that computes a macro-averaged F-score.

  The macro-averaged F-score summarize how well the classifier's `k` predictions
  align with the observed/gold instances of `k`. It additionally cares about
  all the classes equally regardless of their size.

  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 an macro-averaged 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)
    for k in range(initial_category_index, n_categories):
      _, _, _, precision, recall = _precision_recall(predictions, targets, k)
      f_scores = jnp.append(f_scores, _f_score(precision, recall, beta2))
    return jnp.mean(f_scores)

  return base.Fn('MacroAveragedFScore', f)
Returns a layer that computes a macro-averaged F-score.

  The macro-averaged F-score summarize how well the classifier's `k` predictions
  align with the observed/gold instances of `k`. It additionally cares about
  all the classes equally regardless of their size.

  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 an macro-averaged F-score across all the classes.
  
trax/layers/metrics.py