code: def focal_softmax_cross_entropy(
      logits: jnp.ndarray,
      one_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
      label_smoothing: Optional[float] = None,
      label_weights: Optional[jnp.ndarray] = None,
      logits_normalized: bool = False,
      gamma: Optional[float] = 2.0,
      keep_label_dimension: bool = False,
  ) -> jnp.ndarray:
      """Computes focal softmax cross-entropy given logits and targets.
  
      Focal loss as defined in https://arxiv.org/abs/1708.02002. Assuming y is the
      target vector and p is the predicted probability for the class, then:
  
      p_t = p if y == 1 and 1-p otherwise
      Focal loss = -(1-p_t)**gamma * log(p_t)
  
      NOTE: this is weighted unnormalized computation of loss that returns the loss
      of examples in the batch. If you are using it as a loss function, you can
      use the normalilzed version as:
      ```
        unnormalized_loss = focal_softmax_cross_entropy(...)
        if weights is not None:
          normalization = weights.sum()
        else:
          normalization = np.prod(one_hot_targets.shape[:-1])
        loss = jnp.sum(unnormalized_loss) / (normalization + 1e-8)
      ```
  
      Args:
        logits: Output of model in shape [batch, ..., num_classes].
        one_hot_targets: One hot vector of shape [batch, ..., num_classes].
        weights: None or array of shape [batch, ...] (rank of one_hot_targets -1).
        label_smoothing: Scalar to use to smooth the one-hot labels.
        label_weights: Weight per label of shape [num_classes].
        logits_normalized: If True, the logits are assumed to be log probs.
        gamma: Modulating factor of the focal loss.
        keep_label_dimension: If True, the class dimension of the output loss is not
          summed over.
  
      Returns:
        The loss of the examples in the given batch.
      """
      loss = weighted_unnormalized_softmax_cross_entropy(
          logits,
          one_hot_targets,
          weights=None,
          label_smoothing=label_smoothing,
          label_weights=label_weights,
          logits_normalized=logits_normalized,
          keep_label_dimension=True,
      )
      prob = (
          jnp.exp(logits)
          if logits_normalized
          else jax.nn.softmax(logits)
      )
      prob = (prob * one_hot_targets).sum(axis=-1, keepdims=True)
      loss *= (1.0 - prob) ** gamma
      if weights is not None:
          loss = apply_weights(loss, weights)
  
      if not keep_label_dimension:
          loss = loss.sum(axis=-1)
  
      return loss
  
imports: import functools
  from typing import Optional, Any, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 21 for calling a global function with example usage: focal_softmax_cross_entropy(logits, one_hot_targets, weights, label_smoothing, label_weights, logits_normalized, gamma, keep_label_dimension) and returns: loss