class_docstr: Tests for DeformableDETRDecoderLayer.
class_docstr_tok: ["Tests","for","DeformableDETRDecoderLayer","."]
docstr: Test DeformableDETRDecoderLayer output shape.
docstr_tok: ["Test","DeformableDETRDecoderLayer","output","shape","."]
code: def test_decoder_layer_output_shape(self, shapes, num_heads):
      """Test DeformableDETRDecoderLayer output shape."""
      rng = random.PRNGKey(8877)
      bs, len_q, embed_dim = 4, 16, 32
      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, num_levels, 4))
      value = jnp.ones((bs, len_v, 8), dtype=jnp.float32)
      pad_mask = jnp.ones((bs, len_v), dtype=bool)
  
      model = DeformableDETRDecoderLayer(
          spatial_shapes=shapes,
          embed_dim=embed_dim,
          num_heads=num_heads,
          num_levels=num_levels,
          num_reference_points=1,
          ffn_dim=32,
          dropout=0.1,
          compiler_config=compiler_config,
      )
  
      params = model.init(
          rng, query, query_pos, ref_points, value, pad_mask, False
      )
      apply = jax.jit(model.apply, static_argnums=6)
      out = apply(
          params, query, query_pos, ref_points, value, pad_mask, False
      )
      self.assertEqual(out.shape, (bs, len_q, embed_dim))
  
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 1 for calling the DeformableDETRDecoderLayerTest obj's underlying member method code with example usage: obj.test_decoder_layer_output_shape(shapes, num_heads) without return types