class_docstr: Tests for detr model.py.
class_docstr_tok: ["Tests","for","detr","model.py","."]
docstr: Test DETRTransformer output shape.
docstr_tok: ["Test","DETRTransformer","output","shape","."]
code: def test_detr_transformer_output_shape(
      self, return_intermediate, expected_output_shape
  ):
      """Test DETRTransformer output shape."""
      rng = random.PRNGKey(0)
      inputs_shape = (8, 20 * 20, 64)
      num_query_objects = 100
      num_decoder_layers = 6
      inputs = jnp.array(np.random.normal(size=inputs_shape))
      query_pos_emb = jnp.array(
          np.random.normal(size=(1, num_query_objects, 64))
      )
  
      # test output shape of DETR Transformer model
      detr_transformer_def = model.DETRTransformer(
          num_queries=num_query_objects,
          qkv_dim=64,
          num_heads=4,
          num_decoder_layers=num_decoder_layers,
          return_intermediate_dec=return_intermediate,
      )
  
      (outputs, _), _ = detr_transformer_def.init_with_output(
          rng, inputs, query_pos_emb=query_pos_emb
      )
      self.assertEqual(outputs.shape, expected_output_shape)
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  from flax import traverse_util
  from flax.core import unfreeze
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.detr import model
label_desc: Use image node_id 4 for calling the DETRModulesTest obj's underlying member method code with example usage: obj.test_detr_transformer_output_shape(return_intermediate, expected_output_shape) without return types