class_docstr: Layer of DeformableDETR decoder.
  
  Uses MultiScaleDeformableAttention for cross-attention and typical DETR dense
  attention for the self-attention.
  
  Attributes:
    spatial_shapes: (h, w) for each feature level.
    embed_dim: Size of the hidden embedding dimension, used for query, value,
      embeddings, and outputs.
    num_heads: Number of heads.
    num_levels: Number of feature levels.
    num_points: Number of points in deformable attention.
    dropout: Dropout rate.
    ffn_dim: Hidden dimension for feed-forward/MLP network.
    dtype: Data type of the computation (default: float32).
docstr: Compute decoder layer.
  
  Args:
    query: [bs, len_q, embed_dim] of queries.
    query_pos: [bs, len_q, embed_dim] of position embedding for each query
      position.
    ref_points: [bs, len_q, num_levels, box_dim] of reference points for each
      query. box_dim is in {2, 4} as ref_points can either be box cxcy or
      cxcywh.
    value: [bs, len_v, embed_dim] of values to be applied in cross-attention.
    pad_mask: [bs, len_v] of boolean values, where 0 indicates padding in
      value array.
    train: Whether we are in training mode.
  
  Returns:
    [bs, len_q, embed_dim] of decoder-encoded queries.
docstr_tok: ["Compute","decoder","layer",".","Args",":","query",":","[","bs",",","len_q",",","embed_dim","]","of","queries",".","query_pos",":","[","bs",",","len_q",",","embed_dim","]","of","position","embedding","for","each","query","position",".","ref_points",":","[","bs",",","len_q",",","num_levels",",","box_dim","]","of","reference","points","for","each","query",".","box_dim","is","in","{","2",",","4","}","as","ref_points","can","either","be","box","cxcy","or","cxcywh",".","value",":","[","bs",",","len_v",",","embed_dim","]","of","values","to","be","applied","in","cross-attention",".","pad_mask",":","[","bs",",","len_v","]","of","boolean","values",",","where","0","indicates","padding","in","value","array",".","train",":","Whether","we","are","in","training","mode",".","Returns",":","[","bs",",","len_q",",","embed_dim","]","of","decoder-encoded","queries","."]
code: def __call__(
      self,
      query: jnp.ndarray,
      query_pos: jnp.ndarray,
      ref_points: jnp.ndarray,
      value: jnp.ndarray,
      pad_mask: jnp.ndarray,
      train: bool,
  ) -> jnp.ndarray:
      """Compute decoder layer.
  
      Args:
        query: [bs, len_q, embed_dim] of queries.
        query_pos: [bs, len_q, embed_dim] of position embedding for each query
          position.
        ref_points: [bs, len_q, num_levels, box_dim] of reference points for each
          query. box_dim is in {2, 4} as ref_points can either be box cxcy or
          cxcywh.
        value: [bs, len_v, embed_dim] of values to be applied in cross-attention.
        pad_mask: [bs, len_v] of boolean values, where 0 indicates padding in
          value array.
        train: Whether we are in training mode.
  
      Returns:
        [bs, len_q, embed_dim] of decoder-encoded queries.
      """
      x = MultiHeadDotProductAttention(
          name="self_attn",
          # Scaling is necessary to match pytorch official model that combines
          # the kernel into one init with larger fan_in.
          # qkv_kernel_init=jax.nn.initializers.variance_scaling(
          #     0.5, 'fan_avg', 'uniform'),
          num_heads=self.num_heads,
      )(inputs_q=query, pos_emb_q=query_pos, pos_emb_k=query_pos)
  
      query = query + nn.Dropout(rate=self.dropout)(
          x, deterministic=not train
      )
      # TODO(tonysherbondy): Reverse layer norm naming and just fix in
      # torch_param map.
      query = nn.LayerNorm(dtype=self.dtype, name="norm2")(query)
  
      # cross attention
      x = MultiScaleDeformableAttention(
          spatial_shapes=self.spatial_shapes,
          num_levels=self.num_levels,
          num_heads=self.num_heads,
          num_points=self.num_reference_points,
          embed_dim=self.embed_dim,
          compiler_config=self.compiler_config,
          dtype=self.dtype,
          name="cross_attn",
      )(query + query_pos, ref_points, value, pad_mask, train)
  
      query = query + nn.Dropout(rate=self.dropout)(
          x, deterministic=not train
      )
      query = nn.LayerNorm(dtype=self.dtype, name="norm1")(query)
  
      # FeedForward Network.
      # TODO(tonysherbondy): Extract as module since its the same as other FFN.
      bias_range = 1 / np.sqrt(x.shape[-1])
      x = nn.Dense(
          self.ffn_dim,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=uniform_initializer(
              -bias_range, bias_range, dtype=self.dtype
          ),
          name="linear1",
      )(query)
  
      x = nn.relu(x)
      x = nn.Dropout(rate=self.dropout)(x, deterministic=not train)
      bias_range = 1 / np.sqrt(x.shape[-1])
      x = nn.Dense(
          self.embed_dim,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=uniform_initializer(
              -bias_range, bias_range, dtype=self.dtype
          ),
          name="linear2",
      )(x)
      x = nn.Dropout(rate=self.dropout)(x, deterministic=not train)
      query = query + x
      query = nn.LayerNorm(name="norm3")(query)
  
      # TODO(tonysherbondy): Consider clamping values between [-max, max].
      return query
  
imports: import functools
  from typing import Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.deformable_detr.attention import MultiScaleDeformableAttention
  from scenic.projects.baselines.detr.model import MultiHeadDotProductAttention
label_desc: Use image node_id 1 for calling the DeformableDETRDecoderLayer obj's underlying member method code with example usage: obj.__call__(query, query_pos, ref_points, value, pad_mask, train) and returns: query