class_docstr: Tests for DeformableDETREncoder.
class_docstr_tok: ["Tests","for","DeformableDETREncoder","."]
docstr: Test DeformableDETREncoderLayer output shape.
docstr_tok: ["Test","DeformableDETREncoderLayer","output","shape","."]
code: def test_encoder_layer_output_shape(self, num_layers, shapes):
      """Test DeformableDETREncoderLayer output shape."""
      rng = random.PRNGKey(8877)
  
      # Setup size params.
      embed_dim, num_heads = 8, 2
      len_qkv = np.array(shapes).prod(axis=-1).sum()
      bs, num_levels, ref_dim = 2, len(shapes), 2
      src = jnp.array(np.random.normal(size=(bs, len_qkv, embed_dim)))
      ref_points = jnp.array(
          np.random.normal(size=(bs, len_qkv, num_levels, ref_dim))
      )
      pad_mask = np.zeros(src.shape[:-1], dtype=bool)
      pos_embed = jnp.array(np.random.normal(size=src.shape))
  
      # Compute.
      model = DeformableDETREncoder(
          spatial_shapes=shapes,
          embed_dim=embed_dim,
          num_heads=num_heads,
          num_levels=num_levels,
          num_layers=num_layers,
          num_reference_points=1,
          ffn_dim=8,
          dropout=0.1,
          compiler_config=compiler_config,
      )
  
      out, init_params = model.init_with_output(
          rng, src, pos_embed, ref_points, pad_mask, False
      )
      self.assertSequenceEqual(out.shape, (bs, len_qkv, embed_dim))
  
      # Can jit.
      run = jax.jit(model.apply, static_argnums=5)
      out2 = run(
          init_params, src, pos_embed, ref_points, pad_mask, False
      )
      self.assertSequenceEqual(out2.shape, (bs, len_qkv, 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 DeformableDETREncoderTest obj's underlying member method code with example usage: obj.test_encoder_layer_output_shape(num_layers, shapes) without return types