code: def compute_box_bias(
      feature_map: jnp.ndarray,
      padding_mask: Optional[jnp.ndarray] = None,
      kind: str = "both",
  ) -> jnp.ndarray:
      """Computes spatial bias for grid."""
      # The box center is biased to its position on the feature grid:
      xy = normalized_grid_corner_coordinates(feature_map, padding_mask)
      xy = jnp.clip(xy, 0.0, 1.0)
  
      if kind in ["both", "location"]:
          # Unnormalize xy (i.e., apply logit function/sigmoid^-1).
          xy_bias = logit(xy)
      else:
          xy_bias = jnp.zeros_like(xy)
  
      if kind in ["both", "size"]:
          # The box size is biased to the patch size:
          wh_bias = logit(
              jnp.full_like(xy_bias, 1.0 / feature_map.shape[-2])
          )
      else:
          wh_bias = jnp.zeros_like(xy_bias)
  
      return jnp.concatenate([xy_bias, wh_bias], axis=-1)
  
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 4 for calling a global function with example usage: compute_box_bias(feature_map, padding_mask, kind) and returns: jnp