code: def weighted_unnormalized_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,
      keep_label_dimension: bool = False,
  ) -> jnp.ndarray:
      """Computes weighted softmax cross entropy give logits and targets.
  
      This computes sum_(x,y) softmax-ce(x, y) for a single, potentially padded
      minibatch. If the minibatch is padded (that is it contains null examples)
      it is assumed that weights is a binary mask where 0 indicates that the
      example is null.
  
      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 x ...] (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 already be normalized.
        keep_label_dimension: If True, the class dimension of the output loss is not
          summed over.
  
      Returns:
        The softmax cross entropy of the examples in the given batch.
      """
      if logits.ndim != one_hot_targets.ndim:
          raise ValueError(
              "Incorrect shapes. Got shape %s logits and %s one_hot_targets"
              % (str(logits.shape), str(one_hot_targets.shape))
          )
  
      # Optionally apply label smoothing.
      if label_smoothing is not None:
          one_hot_targets = apply_label_smoothing(
              one_hot_targets, label_smoothing
          )
  
      # Optionally apply label weights.
      if label_weights is not None:
          one_hot_targets *= label_weights
  
      if not logits_normalized:
          logits = nn.log_softmax(logits)
      loss = -one_hot_targets * logits
      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 10 for calling a global function with example usage: weighted_unnormalized_softmax_cross_entropy(logits, one_hot_targets, weights, label_smoothing, label_weights, logits_normalized, keep_label_dimension) and returns: loss