code: def bilinear_interpolate(
      im: Array, grid: Array, w: int, h: int
  ) -> jnp.ndarray:
      """Performs 2D bilinear interpolation.
  
      It is assumed that the center of the top-left pixel in `im` has coordinate
      (0.5, 0.5). If you want a different mapping, transform `grid` before calling
      this function. For example, if you want the center of the top-left pixel to
      have coordinate (0, 0), pass `grid + [0.5, 0.5]` into the function instead.
  
      Args:
        im: [image_height * image_width, nembed] a flattened 2D image.
        grid: [..., 2] normalized sampling grid.
        w: Image width.
        h: Image height.
  
      Returns:
        [..., nembed] array of interpolated values.
      """
      im = im.reshape(h, w, -1)
      im = jnp.pad(im, ((1, 1), (1, 1), (0, 0)), "empty")
      im = im.reshape((h + 2) * (w + 2), -1)
  
      x = grid[..., 0] * w
      y = grid[..., 1] * h
      x -= 0.5
      y -= 0.5
      x0 = jnp.floor(x).astype(int)
      x1 = x0 + 1
      y0 = jnp.floor(y).astype(int)
      y1 = y0 + 1
  
      # An important observation that can be made is that we can group the gathering
      # of (x, y), (x + 1, y), (x, y + 1), (x + 1, y + 1) together since they use
      # the same indices for gathering. After packing, we perform one gathering with
      # 4x feature dimension instead of four gatherings with 1x feature dimension,
      # which greatly improves the speed on TPUs, since TPUs have very slow
      # gathering and also forces the feature dimension onto the 128-dimensional
      # sublanes.
  
      # prepare for packing
      indices_y0_offset = (jnp.arange(h + 1) * (w + 2))[:, None]
      indices_y1_offset = (jnp.arange(1, h + 2) * (w + 2))[:, None]
      indices_x0 = jnp.arange(w + 1)
      indices_x1 = jnp.arange(1, w + 2)
      im00 = im[(indices_y0_offset + indices_x0).flatten()]
      im10 = im[(indices_y1_offset + indices_x0).flatten()]
      im01 = im[(indices_y0_offset + indices_x1).flatten()]
      im11 = im[(indices_y1_offset + indices_x1).flatten()]
  
      # pack: (nembed * 4, (h + 1) * (w + 1))
      im_packed = jnp.concatenate([im00, im10, im01, im11], axis=-1)
  
      # gather
      indices11 = jnp.clip(y1, 0, h + 1) * (w + 1) + jnp.clip(
          x1, 0, w + 1
      )
      im_gathered = im_packed[indices11]
  
      # unpack
      im_a, im_b, im_c, im_d = jnp.split(im_gathered, 4, axis=-1)
  
      # Mark indices out-of-bounds.
      x0_out = jnp.logical_or(x0 < 0, x0 > w - 1)
      y0_out = jnp.logical_or(y0 < 0, y0 > h - 1)
      x1_out = jnp.logical_or(x1 < 0, x1 > w - 1)
      y1_out = jnp.logical_or(y1 < 0, y1 > h - 1)
      out00 = jnp.logical_or(x0_out, y0_out)
      out01 = jnp.logical_or(x0_out, y1_out)
      out10 = jnp.logical_or(x1_out, y0_out)
      out11 = jnp.logical_or(x1_out, y1_out)
  
      # Set weights where weights for out-of-bound pixels are forced to be 0.
      wa = jnp.where(out00, 0, (x1 - x) * (y1 - y))
      wb = jnp.where(out01, 0, (x1 - x) * (y - y0))
      wc = jnp.where(out10, 0, (x - x0) * (y1 - y))
      wd = jnp.where(out11, 0, (x - x0) * (y - y0))
  
      return (
          jnp.einsum("...e,...->...e", im_a, wa)
          + jnp.einsum("...e,...->...e", im_b, wb)
          + jnp.einsum("...e,...->...e", im_c, wc)
          + jnp.einsum("...e,...->...e", im_d, wd)
      )
  
file_path: scenic/scenic/projects/baselines/deformable_detr/attention.py
imports: import functools
  import math
  from typing import Callable, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.core import Shape
  import jax.numpy
  import ml_collections
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: bilinear_interpolate(im, grid, w, h) and returns: unknown