code: def deform_attn_sampling_fn(
      values: Array,
      sampling_locations: Array,
      attn_weights: Array,
      shapes: Tuple[Tuple[int, int], ...],
      use_remat: bool,
      mode: str,
  ) -> jnp.ndarray:
      """Performs deformable attention sampling calculation.
  
      Args:
        values: [bs, len_v, nembed]-array of values.
        sampling_locations: [bs, nlevels, npoints * len_q, 2]-array of
          sampling locations.
        attn_weights: [bs, len_q, nlevels * npoints]-array of attention weights.
        shapes: Static tuple of image shapes for each level to unflatten len_v.
        use_remat: Flag for rematerialization.
        mode: Determines how batching is performed, can be one of the following:
          'loop', 'map', 'vmap', 'auto'.
  
      Returns:
        [bs, len_q, nembed]-array
      """
      nembed = values.shape[-1]
      len_q = attn_weights.shape[1]
      split_indices = np.cumsum(np.array([h * w for h, w in shapes]))[
          :-1
      ]
      # Split values by level.
      values_by_level = jnp.split(values, split_indices, axis=-2)
      sampled_values_all_levels = []
  
      if use_remat:
          attention_fn = jax.remat(
              bilinear_interpolate,
              policy=jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims,
              static_argnums=(2, 3),
          )
      else:
          attention_fn = bilinear_interpolate
  
      def attention_at_index_fn(i, level_idx):
          return attention_fn(
              values_by_level[level_idx][i],
              sampling_locations[:, level_idx][i],
              shapes[level_idx][1],
              shapes[level_idx][0],
          )
  
      def reshape_attn(fn):
          return lambda *args: jnp.reshape(
              fn(*args), (-1, len_q, nembed)
          )
  
      for level_idx in range(len(shapes)):
          fn = functools.partial(
              attention_at_index_fn, level_idx=level_idx
          )
          fn = _map(reshape_attn(fn), mode)
          sampled_values_all_levels.append(
              fn(jnp.array(range(values.shape[0])))
          )
  
      # (bs, nlevels * npoints, len_q, nembed)
      sampled_values_all_levels = jnp.concatenate(
          sampled_values_all_levels, axis=1
      )
  
      # (bs, len_q, nembed)
      return jnp.einsum(
          "blqe,bql->bqe", sampled_values_all_levels, attn_weights
      )
  
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 3 for calling a global function with example usage: deform_attn_sampling_fn(values, sampling_locations, attn_weights, shapes, use_remat, mode) and returns: jnp, attention_fn, function