code: def num_examples(
      logits: jnp.ndarray,
      one_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> Union[jnp.ndarray, int]:
      del logits
      if weights is None:
          return one_hot_targets.shape[0]
      return weights.sum()
  
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 2 for calling a global function with example usage: num_examples(logits, one_hot_targets, weights) and returns: weights, one_hot_targets