class_docstr: Tests for DeformableDETRDecoder.
class_docstr_tok: ["Tests","for","DeformableDETRDecoder","."]
docstr: Test DeformableDETRDecoder output shape.
docstr_tok: ["Test","DeformableDETRDecoder","output","shape","."]
code: def test_decoder_output_shape(self, shapes, ref_dim):
      """Test DeformableDETRDecoder output shape."""
      rng = random.PRNGKey(8877)
      bs, len_q, embed_dim, num_heads = 3, 16, 32, 4
      num_layers = 5
      len_v = np.array(shapes).prod(1).sum()
      num_levels = len(shapes)
  
      query = jnp.ones((bs, len_q, embed_dim), dtype=jnp.float32)
      query_pos = jnp.ones_like(query)
      ref_points = jnp.zeros((bs, len_q, ref_dim))
      value = jnp.ones((bs, len_v, 8), dtype=jnp.float32)
      pad_mask = jnp.ones((bs, len_v), dtype=bool)
      valid_ratios = jnp.ones((bs, num_levels, 2), dtype=jnp.float32)
  
      bbox_embeds = [
          BBoxCoordPredictor(
              mlp_dim=embed_dim, num_layers=3, use_sigmoid=False
          )
          for _ in range(num_layers)
      ]
  
      model = DeformableDETRDecoder(
          spatial_shapes=shapes,
          embed_dim=embed_dim,
          num_heads=num_heads,
          num_layers=num_layers,
          num_levels=num_levels,
          num_reference_points=1,
          bbox_embeds=bbox_embeds,
          ffn_dim=16,
          dropout=0.1,
          compiler_config=compiler_config,
      )
  
      dec_input = dict(
          query=query,
          query_pos=query_pos,
          ref_points=ref_points,
          value=value,
          valid_ratios=valid_ratios,
          pad_mask=pad_mask,
          train=False,
      )
  
      params = model.init(rng, **dec_input)
      apply = jax.jit(model.apply, static_argnames="train")
      out, out_boxes = apply(params, **dec_input)
      self.assertEqual(out.shape, (num_layers, bs, len_q, embed_dim))
      self.assertEqual(out_boxes.shape, (num_layers, bs, len_q, 4))
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.deformable_detr.deformable_transformer import BBoxCoordPredictor
  from scenic.projects.baselines.deformable_detr.deformable_transformer import DeformableDETRDecoder
  from scenic.projects.baselines.deformable_detr.deformable_transformer import DeformableDETRDecoderLayer
  from scenic.projects.baselines.deformable_detr.deformable_transformer import DeformableDETREncoder
  from scenic.projects.baselines.deformable_detr.deformable_transformer import DeformableDETREncoderLayer
  from scenic.projects.baselines.deformable_detr.deformable_transformer import DeformableDETRTransformer
  from scenic.projects.baselines.deformable_detr.deformable_transformer import get_encoder_reference_points
  from scenic.projects.baselines.deformable_detr.deformable_transformer import get_mask_valid_ratio
  from scenic.projects.baselines.deformable_detr.deformable_transformer import inverse_sigmoid
label_desc: Use image node_id 2 for calling the DeformableDETRDecoderTest obj's underlying member method code with example usage: obj.test_decoder_output_shape(shapes, ref_dim) without return types