code: def WeightedFScore(beta=1.0, 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)
  
imports: from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
label_desc: Use image node_id 19 for calling a global function with example usage: WeightedFScore(beta, initial_category_index) and returns: base, jnp