class_docstr: Layer for MultiScaleDeformableAttention.
class_docstr_tok: ["Layer","for","MultiScaleDeformableAttention","."]
docstr: Calculates multi-scale multi-head deformable attention.
  
  Args:
    query: [bs, len_q, embed_dim]-ndarray of queries.
    ref_points: [bs, len_q, num_levels, box_dim]-ndarray of reference points
      for each query.
    value: [bs, len_v, embed_dim]-ndarray of values.
    pad_mask: [bs, len_v]-ndarray of boolean values, where 0 indicates pad.
    train: Whether we are in training mode.
  
  Returns:
    Attention weighted values based on the queries and reference points
    relevant to the queries.
docstr_tok: ["Calculates","multi-scale","multi-head","deformable","attention",".","Args",":","query",":","[","bs",",","len_q",",","embed_dim","]","-ndarray","of","queries",".","ref_points",":","[","bs",",","len_q",",","num_levels",",","box_dim","]","-ndarray","of","reference","points","for","each","query",".","value",":","[","bs",",","len_v",",","embed_dim","]","-ndarray","of","values",".","pad_mask",":","[","bs",",","len_v","]","-ndarray","of","boolean","values",",","where","0","indicates","pad",".","train",":","Whether","we","are","in","training","mode",".","Returns",":","Attention","weighted","values","based","on","the","queries","and","reference","points","relevant","to","the","queries","."]
code: def __call__(
      self,
      query: jnp.ndarray,
      ref_points: jnp.ndarray,
      value: jnp.ndarray,
      pad_mask: jnp.ndarray,
      train: bool,
  ) -> jnp.ndarray:
      """Calculates multi-scale multi-head deformable attention.
  
      Args:
        query: [bs, len_q, embed_dim]-ndarray of queries.
        ref_points: [bs, len_q, num_levels, box_dim]-ndarray of reference points
          for each query.
        value: [bs, len_v, embed_dim]-ndarray of values.
        pad_mask: [bs, len_v]-ndarray of boolean values, where 0 indicates pad.
        train: Whether we are in training mode.
  
      Returns:
        Attention weighted values based on the queries and reference points
        relevant to the queries.
      """
      assert (
          self.embed_dim % self.num_heads == 0
      ), "`embed_dim` must be divisible by `num_heads`."
      ref_dim = ref_points.shape[-1]
      bs, len_q, _ = query.shape
      _, len_v, _ = value.shape
  
      nembed = self.embed_dim // self.num_heads
  
      # (bs, len_v, nheads, nembed)
      value = nn.DenseGeneral(
          features=(self.num_heads, nembed),
          param_dtype=self.dtype,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=nn.initializers.zeros,
          name="value_proj",
      )(value)
      # (bs, len_v, nheads, nembed)
      value = jnp.where(pad_mask[..., None, None], value, 0)
      # (bs, nheads, len_v, nembed)
      value = value.transpose(0, 2, 1, 3)
      # (bs * nheads, len_v, nembed)
      value = value.reshape(bs * self.num_heads, len_v, -1)
  
      # (bs, len_q, nheads, nlevels, npoints, 2)
      sampling_offsets = nn.DenseGeneral(
          features=(
              self.num_heads,
              self.num_levels,
              self.num_points,
              2,
          ),
          param_dtype=self.dtype,
          kernel_init=nn.initializers.zeros,
          bias_init=self.pos_grid_init,
          name="sampling_offsets",
      )(query)
  
      if ref_dim == 2:
          # (bs, len_q, 1, nlevels, 1, 2)
          ref_xy = ref_points[:, :, None, :, None, :]
          # (bs, len_q, nheads, nlevels, npoints, 2)
          normalized_offsets = sampling_offsets / self.offset_norm
      elif ref_dim >= 4:
          # (bs, len_q, 1, nlevels, 1, 2)
          ref_xy = ref_points[:, :, None, :, None, :2]
          # (bs, len_q, 1, nlevels, 1, 2)
          ref_wh = ref_points[:, :, None, :, None, 2:4]
          # (bs, len_q, nheads, nlevels, npoints, 2)
          normalized_offsets = (
              sampling_offsets / (2 * self.num_points) * ref_wh
          )
  
      # (bs, len_q, nheads, nlevels, npoints, 2)
      sampling_locations = ref_xy + normalized_offsets
      # (bs, nheads, nlevels, npoints, len_q, 2)
      sampling_locations = sampling_locations.transpose(
          0, 2, 3, 4, 1, 5
      )
      # (bs * nheads, nlevels, npoints * len_q, 2)
      sampling_locations = sampling_locations.reshape(
          bs * self.num_heads,
          self.num_levels,
          self.num_points * len_q,
          2,
      )
  
      # (bs, len_q, nheads, nlevels * npoints)
      attn_weights = nn.DenseGeneral(
          features=(self.num_heads, self.num_levels * self.num_points),
          param_dtype=self.dtype,
          kernel_init=nn.initializers.zeros,
          bias_init=nn.initializers.zeros,
          name="attention_weights",
      )(query)
      # (bs, len_q, nheads, nlevels * npoints)
      attn_weights = nn.softmax(attn_weights)
      # (bs, nheads, len_q, nlevels * npoints)
      attn_weights = attn_weights.transpose(0, 2, 1, 3)
      # (bs * nheads, len_q, nlevels * npoints)
      attn_weights = attn_weights.reshape(
          bs * self.num_heads, len_q, self.num_levels * self.num_points
      )
      if train:
          use_remat = self.compiler_config.train_remat
      else:
          use_remat = False
      # (bs * nheads, len_q, nembed)
      x = deform_attn_sampling_fn(
          values=value,
          sampling_locations=sampling_locations,
          attn_weights=attn_weights,
          shapes=self.spatial_shapes,
          use_remat=use_remat,
          mode=self.compiler_config.attention_batching_mode,
      )
  
      # (bs, nheads, len_q, nembed)
      x = x.reshape(bs, self.num_heads, len_q, nembed)
      # (bs, len_q, embed_dim)
      return nn.DenseGeneral(
          features=self.embed_dim,
          axis=(-3, -1),
          param_dtype=self.dtype,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=nn.initializers.zeros,
          name="output_proj",
      )(x)
  
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 the MultiScaleDeformableAttention obj's underlying member method code with example usage: obj.__call__(query, ref_points, value, pad_mask, train) and returns: nn