class_docstr: Tests for DeformableDETRTransformer and utilities.
class_docstr_tok: ["Tests","for","DeformableDETRTransformer","and","utilities","."]
docstr: Test DeformableDETRTransformer output shape.
docstr_tok: ["Test","DeformableDETRTransformer","output","shape","."]
code: def test_transformer_output_shape(self, shapes, num_dec_layers):
      """Test DeformableDETRTransformer output shape."""
      rng = random.PRNGKey(8877)
  
      bs, num_heads, num_queries = 2, 3, 5
      embed_dim = 2 * num_heads
      inputs = []
      pad_masks = []
      pos_embeds = []
      for h, w in shapes:
          inputs.append(
              jnp.ones((bs, h, w, embed_dim), dtype=jnp.float32)
          )
          pad_masks.append(jnp.ones((bs, h, w), dtype=jnp.bool_))
          pos_embeds.append(
              jnp.ones((bs, h * w, embed_dim), dtype=jnp.float32)
          )
  
      bbox_embeds = [
          BBoxCoordPredictor(
              mlp_dim=embed_dim, num_layers=3, use_sigmoid=False
          )
          for _ in range(num_dec_layers)
      ]
  
      # Compute.
      model = DeformableDETRTransformer(
          enc_embed_dim=embed_dim,
          embed_dim=embed_dim,
          num_queries=num_queries,
          num_heads=num_heads,
          num_dec_layers=num_dec_layers,
          num_enc_layers=2,
          ffn_dim=8,
          bbox_embeds=bbox_embeds,
          num_enc_points=1,
          num_dec_points=2,
          dropout=0.1,
          compiler_config=compiler_config,
      )
  
      (
          out,
          ref_boxes,
          init_ref_points,
      ), init_params = model.init_with_output(
          rng, inputs, pad_masks, pos_embeds, False
      )
  
      self.assertSequenceEqual(
          out.shape, (num_dec_layers, bs, num_queries, embed_dim)
      )
      self.assertSequenceEqual(
          ref_boxes.shape, (num_dec_layers, bs, num_queries, 4)
      )
      self.assertSequenceEqual(
          init_ref_points.shape, (bs, num_queries, 2)
      )
  
      # Can jit.
      run = jax.jit(model.apply, static_argnums=4)
      out, ref_boxes, init_ref_points = run(
          init_params, inputs, pad_masks, pos_embeds, False
      )
      self.assertSequenceEqual(
          out.shape, (num_dec_layers, bs, num_queries, embed_dim)
      )
      self.assertSequenceEqual(
          ref_boxes.shape, (num_dec_layers, bs, num_queries, 4)
      )
      self.assertSequenceEqual(
          init_ref_points.shape, (bs, num_queries, 2)
      )
  
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 3 for calling the DeformableDETRTransformerTest obj's underlying member method code with example usage: obj.test_transformer_output_shape(shapes, num_dec_layers) without return types