code: def ponder_loss_fn(
      all_p: jnp.ndarray,
      p_g: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Ponder Loss for PonderNet.
  
      Args:
        all_p: Input array of any shape.
        p_g: Input array of any shape.
        weights: None or array of any shape.
  
      Returns:
        loss: A scaler to regularize the PonderNet.
      """
      all_p = all_p.transpose((1, 0))
      if weights is not None:
          normalization = weights.sum()
      else:
          normalization = np.prod(all_p.shape[0])
      p_g = jnp.expand_dims(p_g, axis=0)
      p_g = p_g.repeat(all_p.shape[0], axis=0)
      # Calculate the KL div between all_p and p_g
      loss = jnp.sum(
          p_g * (jnp.log(p_g + 1e-8) - jnp.log(all_p + 1e-8))
      ) / (normalization + 1e-8)
      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.pondernet import layers
label_desc: Use image node_id 1 for calling a global function with example usage: ponder_loss_fn(all_p, p_g, weights) and returns: loss