def CategoryCrossEntropy(label_smoothing=None):
  r"""Returns a layer that computes cross-entropy from activations and integers.

  The layer takes two inputs:

    - A batch of activation vectors. The components in a given vector should
      be pre-softmax activations (mappable to a probability distribution via
      softmax). For performance reasons, the softmax and cross-entropy
      computations are combined inside the layer.

    - A batch of target categories; each target is an integer in
      :math:`\{0, ..., N-1\}`, where :math:`N` is the activation vector
      depth/dimensionality.

  To compute cross-entropy per batch item, the layer derives probability
  distributions:

    - from model output (vectors): :math:`\ q = \text{softmax}(v)`

    - from target categories (integers): :math:`\ p = \text{one_hot}(n)` or
      :math:`p = (1-\varepsilon)\cdot\text{one_hot}(n) + \frac{\varepsilon}{N}`,
      where :math:`\varepsilon` is the label smoothing factor.

  (The conversion of integer category targets to one-hot vectors amounts to
  assigning all the probability mass to the target category.) Cross-entropy
  per batch item is computed between the resulting distributions:

  .. math::
      \text{cross_entropy} = - \sum_{i=0}^{N-1} p_i \log q_i

  The layer returns the average of these cross-entropy values over all items in
  the batch.

  Args:
    label_smoothing: Creates soft targets if provided. Must be between 0 and 1.
  """
  def f(model_output, targets):  # pylint: disable=invalid-name
    cross_entropies = _category_cross_entropy(
        model_output, targets, label_smoothing, 0.0)
    return jnp.average(cross_entropies)

  return base.Fn('CategoryCrossEntropy', f)
Returns a layer that computes cross-entropy from activations and integers.

  The layer takes two inputs:

    - A batch of activation vectors. The components in a given vector should
      be pre-softmax activations (mappable to a probability distribution via
      softmax). For performance reasons, the softmax and cross-entropy
      computations are combined inside the layer.

    - A batch of target categories; each target is an integer in
      :math:`\{0, ..., N-1\}`, where :math:`N` is the activation vector
      depth/dimensionality.

  To compute cross-entropy per batch item, the layer derives probability
  distributions:

    - from model output (vectors): :math:`\ q = \text{softmax}(v)`

    - from target categories (integers): :math:`\ p = \text{one_hot}(n)` or
      :math:`p = (1-\varepsilon)\cdot\text{one_hot}(n) + \frac{\varepsilon}{N}`,
      where :math:`\varepsilon` is the label smoothing factor.

  (The conversion of integer category targets to one-hot vectors amounts to
  assigning all the probability mass to the target category.) Cross-entropy
  per batch item is computed between the resulting distributions:

  .. math::
      \text{cross_entropy} = - \sum_{i=0}^{N-1} p_i \log q_i

  The layer returns the average of these cross-entropy values over all items in
  the batch.

  Args:
    label_smoothing: Creates soft targets if provided. Must be between 0 and 1.
  
trax/layers/metrics.py