code: def ponder_loss_fn(
      ponder_times: jnp.ndarray,
      remainders: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Ponder Loss for UT.
  
      Args:
        ponder_times: Input array of any shape.
        remainders: Input array of any shape.
        weights: None or array of any shape.
  
      Returns:
        loss: A scaler to regularize the ACT
      """
  
      if weights is not None:
          normalization = weights.sum() + 1e-8
      else:
          normalization = np.prod(ponder_times.shape)
  
      p_t = ponder_times + remainders
      # We only do mean when ac_config.act_level == per_example.
      # When ac_config.act_level == per_token, we do sum first.
      if jnp.ndim(p_t) == 1:
          loss = jnp.sum(p_t) / normalization
      else:
          loss = jnp.sum(jnp.sum(p_t, axis=-1)) / normalization
      return loss
  
imports: from typing import Any, Optional
  import flax.linen
  from flax.training import common_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.baselines.universal_transformer import layers
label_desc: Use image node_id 1 for calling a global function with example usage: ponder_loss_fn(ponder_times, remainders, weights) and returns: loss