code: def point_count(
      logits: jnp.ndarray,
      one_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> float:
      """Computes number of pixels 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) pixels in the input.
      """
      del logits
      if weights is None:
          return np.prod(one_hot_targets.shape[:3])
      assert (
          weights.ndim == 2
      ), "For segmentation task, the weights should be a point level mask."
      return weights.sum()
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  import flax.linen
  from flax.training import common_utils
  from immutabledict import immutabledict
  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.segmentation_model import SegmentationModel
  from scenic.projects.pointcloud.models import PointCloudTransformerEncoder
label_desc: Use image node_id 1 for calling a global function with example usage: point_count(logits, one_hot_targets, weights) and returns: weights, np