code: def weighted_correctly_classified(
      logits: jnp.ndarray,
      one_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Computes weighted number of correctly classified over the given batch.
  
      This computes the weighted number of correctly classified examples/pixels in a
      single, potentially padded minibatch. If the minibatch/inputs is padded (i.e.,
      it contains null examples/pad pixels) it is assumed that weights is a binary
      mask where 0 indicates that the example/pixel is null/padded. We assume the
      trainer will aggregate and divide by number of samples.
  
      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).
  
      Returns:
        The number of correctly classified 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))
          )
      preds = jnp.argmax(logits, axis=-1)
      targets = jnp.argmax(one_hot_targets, axis=-1)
      correct = jnp.equal(preds, targets)
  
      if weights is not None:
          correct = apply_weights(correct, weights)
  
      return correct.astype(jnp.int32)
  
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 4 for calling a global function with example usage: weighted_correctly_classified(logits, one_hot_targets, weights) and returns: correct