code: def normalized_grid_corner_coordinates(
      feature_map: jnp.ndarray, padding_mask: Optional[jnp.ndarray]
  ) -> jnp.ndarray:
      """Computes normalized xy corner coords from feature_map or padding_mask."""
      # Note 1: it computes not the centers of grid patches, but the patch corner
      # coordinates (for a grid patch from 0 to 0.1, it returns 0.1 not 0.05).
      # Note 2: behavior is quite different for feature_map and padding_mask inputs.
      if padding_mask is None:
          assert feature_map.ndim == 4  # [B, H, W, C]
          h, w = feature_map.shape[1:3]
          xy = np.stack(
              np.meshgrid(np.arange(1, w + 1), np.arange(1, h + 1)),
              axis=-1,
          ).astype(np.float32)
          xy /= np.array([w, h], np.float32)
      else:
          assert padding_mask.ndim == 3  # [B, H, W]
          y = jnp.cumsum(padding_mask, axis=1)
          x = jnp.cumsum(padding_mask, axis=2)
          xy = jnp.stack(
              [x / (x[:, :, -1:] + 1e-6), y / (y[:, -1:] + 1e-6)],
              axis=-1,
          )
      # Flatten h, w dimensions.
      return xy.reshape(*(xy.shape[:-3] + (-1, 2)))
  
imports: from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, Union
  from absl import logging
  from clu import preprocess_spec
  from flax import core
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib import train_utils
  import scipy
label_desc: Use image node_id 3 for calling a global function with example usage: normalized_grid_corner_coordinates(feature_map, padding_mask) and returns: xy