code: def ponder_loss_fn(
      loss_atr: jnp.ndarray, weights: Optional[jnp.ndarray] = None
  ) -> jnp.ndarray:
      """Ponder Loss for UT.
  
      Args:
        loss_atr: 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(loss_atr.shape)
      loss = jnp.sum(loss_atr) / normalization
      return loss
  
imports: from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  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 nn_layers
  from scenic.projects.adatape import layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling a global function with example usage: ponder_loss_fn(loss_atr, weights) and returns: loss