class_docstr: Layer of DETR encoder.
docstr: Single encoder layer using MultiScaleDeformableAttention.
  
  Args:
    src: [bs, len_qkv, embed_dim]-ndarray of values. This is self-attention so
      these will also be used as queries after position embedding added.
    pos_embed: [bs, len_qkv, embed_dim]-ndarray of position embedding for each
      src position.
    ref_points: [bs, len_qkv, num_levels, box_dim]-ndarray of reference points
      for each query. box_dim is in {2, 4} as ref_points can either be box
      cxcy or cxcywh.
    pad_mask: [bs, len_qkv]-ndarray of boolean values, where 1 indicates pad.
    train: Whether we are in training mode.
  
  Returns:
    [bs, len_qkv, embed_dim]-ndarray of encoding from layer.
docstr_tok: ["Single","encoder","layer","using","MultiScaleDeformableAttention",".","Args",":","src",":","[","bs",",","len_qkv",",","embed_dim","]","-ndarray","of","values",".","This","is","self-attention","so","these","will","also","be","used","as","queries","after","position","embedding","added",".","pos_embed",":","[","bs",",","len_qkv",",","embed_dim","]","-ndarray","of","position","embedding","for","each","src","position",".","ref_points",":","[","bs",",","len_qkv",",","num_levels",",","box_dim","]","-ndarray","of","reference","points","for","each","query",".","box_dim","is","in","{","2",",","4","}","as","ref_points","can","either","be","box","cxcy","or","cxcywh",".","pad_mask",":","[","bs",",","len_qkv","]","-ndarray","of","boolean","values",",","where","1","indicates","pad",".","train",":","Whether","we","are","in","training","mode",".","Returns",":","[","bs",",","len_qkv",",","embed_dim","]","-ndarray","of","encoding","from","layer","."]
code: def __call__(
      self,
      src: jnp.ndarray,
      pos_embed: jnp.ndarray,
      ref_points: jnp.ndarray,
      pad_mask: jnp.ndarray,
      train: bool,
  ) -> jnp.ndarray:
      """Single encoder layer using MultiScaleDeformableAttention.
  
      Args:
        src: [bs, len_qkv, embed_dim]-ndarray of values. This is self-attention so
          these will also be used as queries after position embedding added.
        pos_embed: [bs, len_qkv, embed_dim]-ndarray of position embedding for each
          src position.
        ref_points: [bs, len_qkv, num_levels, box_dim]-ndarray of reference points
          for each query. box_dim is in {2, 4} as ref_points can either be box
          cxcy or cxcywh.
        pad_mask: [bs, len_qkv]-ndarray of boolean values, where 1 indicates pad.
        train: Whether we are in training mode.
  
      Returns:
        [bs, len_qkv, embed_dim]-ndarray of encoding from layer.
      """
      query = src + pos_embed
      x = MultiScaleDeformableAttention(
          spatial_shapes=self.spatial_shapes,
          embed_dim=self.embed_dim,
          num_heads=self.num_heads,
          num_levels=self.num_levels,
          num_points=self.num_reference_points,
          compiler_config=self.compiler_config,
          dtype=self.dtype,
          name="self_attn",
      )(query, ref_points, src, pad_mask, train)
  
      x = nn.Dropout(rate=self.dropout)(x, deterministic=not train)
      x = src + x
      x = nn.LayerNorm(name="norm1")(x)
  
      # FeedForward Network.
      residual = x
      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",
      )(x)
      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)
      x = residual + x
      x = nn.LayerNorm(name="norm2")(x)
  
      # TODO(tonysherbondy): Consider clamping values between [-max, max].
      return x
  
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 DeformableDETREncoderLayer obj's underlying member method code with example usage: obj.__call__(src, pos_embed, ref_points, pad_mask, train) and returns: x