code: def num_tokens(
      logits: jnp.ndarray,
      one_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> float:
      """Computes number of tokens in the target to be used for normalization.
  
      It needs to have the same API as other defined metrics.
  
      Args:
        logits: Unused.
        one_hot_targets: Targets, in form of one-hot vectors.
        weights: Input weights (can be used for accounting the padding in the
          input).
  
      Returns:
        Number of (non-padded) tokens in the target.
      """
      del logits
  
      if weights is None:
          return np.prod(one_hot_targets.shape[:2])
      assert (
          weights.ndim == 2
      ), "Weights should be a token level mask of shape [bs, len]."
      return weights.sum()
  
file_path: scenic/scenic/model_lib/base_models/encoder_decoder_model.py
imports: import functools
  from typing import Dict, Optional, Tuple, Union
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 1 for calling a global function with example usage: num_tokens(logits, one_hot_targets, weights) and returns: weights, np