code: def focal_sigmoid_cross_entropy(
      logits: jnp.ndarray,
      multi_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
      label_smoothing: Optional[float] = None,
      label_weights: Optional[jnp.ndarray] = None,
      logits_normalized: bool = False,
      alpha: Optional[float] = 0.5,
      gamma: Optional[float] = 2.0,
  ) -> 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
      alpha_t = alpha if y == 1 and 1-alpha otherwise
  
      Focal loss = -alpha_t * (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_sigmoid_cross_entropy(...)
        if weights is not None:
          normalization = weights.sum()
        else:
          normalization = np.prod(multi_hot_targets.shape[:-1])
        loss = jnp.sum(unnormalized_loss) / (normalization + 1e-8)
      ```
  
      Args:
        logits: Output of model in shape [batch, ..., num_classes].
        multi_hot_targets: Multi-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.
        alpha: Balancing factor of the focal loss.
        gamma: Modulating factor of the focal loss.
  
      Returns:
        The loss of the examples in the given batch.
      """
      # Optionally apply label smoothing.
      if label_smoothing is not None:
          multi_hot_targets = apply_label_smoothing(
              multi_hot_targets, label_smoothing
          )
      if logits_normalized:
          log_p, prob = logits, jnp.exp(logits)
          log_not_p = jnp.log((1 + 1e-6) - prob)
      else:
          log_p, log_not_p = jax.nn.log_sigmoid(
              logits
          ), jax.nn.log_sigmoid(-logits)
  
      loss = -(
          multi_hot_targets * log_p
          + (1.0 - multi_hot_targets) * log_not_p
      )
  
      p_t = jnp.exp(-loss)
      loss *= (1 - p_t) ** gamma
      loss *= alpha * multi_hot_targets + (1 - alpha) * (
          1 - multi_hot_targets
      )
  
      if label_weights is not None:
          loss = loss * label_weights
  
      if weights is not None:
          loss = apply_weights(loss, weights)
      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 22 for calling a global function with example usage: focal_sigmoid_cross_entropy(logits, multi_hot_targets, weights, label_smoothing, label_weights, logits_normalized, alpha, gamma) and returns: loss