code: def weighted_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,
  ) -> jnp.ndarray:
      """Same as weighted_unnormalized, but additionally takes a mean.
  
      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: float scalar to use to smooth the one-hot labels.
        label_weights: Weight per label of shape [num_classes].
  
      Returns:
        The mean cross entropy of the examples in the given batch as a scalar.
      """
      if weights is not None:
          normalization = weights.sum()
      else:
          normalization = np.prod(one_hot_targets.shape[:-1])
  
      unnormalized_softmax_ce = (
          weighted_unnormalized_softmax_cross_entropy(
              logits,
              one_hot_targets,
              weights,
              label_smoothing,
              label_weights,
          )
      )
      return jnp.sum(unnormalized_softmax_ce) / (normalization + 1e-8)
  
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 12 for calling a global function with example usage: weighted_softmax_cross_entropy(logits, one_hot_targets, weights, label_smoothing, label_weights) and returns: unknown