code: def roi_align(
      feature_map: jnp.ndarray, box: jnp.ndarray, output_width: int
  ) -> jnp.ndarray:
      """Extracts a fixed-size feature map for an ROI from a larger feature map.
  
      See the Mask-RCNN paper (https://arxiv.org/abs/1703.06870) for details on
      ROIAlign.
  
      Args:
        feature_map: [H, W, C] map of features from which to crop a region of
          interest.
        box: [cx, cy, w, h] bounding box defining the region of interest.
        output_width: The output region will be resized to [width, width].
  
      Returns:
        Crop of size [width, width] taken from feature_map.
      """
      input_height, input_width, c = feature_map.shape
      output_height = output_width
  
      cx, cy, w, h = jnp.split(box, 4, axis=-1)
      x0 = cx - w / 2
      y0 = cy - h / 2
      w = jnp.maximum(w, 1e-6)
      h = jnp.maximum(h, 1e-6)
      x_scale = output_width / (w * input_width)
      y_scale = output_height / (h * input_height)
  
      return jax.image.scale_and_translate(
          feature_map,
          shape=(output_height, output_width, c),
          spatial_dims=(0, 1),
          scale=jnp.concatenate((y_scale, x_scale)),
          translation=jnp.concatenate(
              (-y0 * output_height / h, -x0 * output_width / w)
          ),
          method="linear",
          precision=jax.lax.Precision("fastest"),
      )
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Optional, Tuple, Union, Sequence
  from absl import logging
  from big_vision.models import bit
  from big_vision.models import vit
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import box_utils
  from scenic.projects.owl_vit import utils
  from scenic.projects.owl_vit.clip import layers
  from scenic.projects.owl_vit.clip import model
label_desc: Use image node_id 2 for calling a global function with example usage: roi_align(feature_map, box, output_width) and returns: jax