class_docstr: Tests for MultiScaleDeformableAttention.
class_docstr_tok: ["Tests","for","MultiScaleDeformableAttention","."]
docstr: Test MultiScaleDeformableAttention output shape.
docstr_tok: ["Test","MultiScaleDeformableAttention","output","shape","."]
code: def test_ms_deformable_attn_output_shape(
      self, ref_dim, shapes, embed_dim, num_heads
  ):
      """Test MultiScaleDeformableAttention output shape."""
      rng = random.PRNGKey(8877)
      # Setup remaining size params.
      bs, len_q, num_points, num_levels = 2, 10, 1, len(shapes)
      len_v = np.array(shapes).prod(axis=-1).sum()
  
      # Set inputs.
      query = jnp.array(np.random.normal(size=(bs, len_q, embed_dim)))
      ref_points = jnp.array(
          np.random.normal(size=(bs, len_q, num_levels, ref_dim))
      )
      value = jnp.array(np.random.normal(size=(bs, len_v, embed_dim)))
      pad_mask = jnp.zeros(value.shape[:-1], dtype=bool)
  
      # Compute.
      model = MultiScaleDeformableAttention(
          embed_dim=embed_dim,
          num_levels=num_levels,
          num_heads=num_heads,
          num_points=num_points,
          spatial_shapes=shapes,
          compiler_config=ml_collections.ConfigDict(
              dict(
                  train_remat=True,
                  eval_remat=False,
                  attention_batching_mode="auto",
              )
          ),
          dtype=jnp.float32,
      )
  
      out, init_params = model.init_with_output(
          rng, query, ref_points, value, pad_mask, True
      )
      self.assertSequenceEqual(out.shape, (bs, len_q, embed_dim))
  
      # Can jit.
      run = jax.jit(model.apply, static_argnums=5)
      out2 = run(init_params, query, ref_points, value, pad_mask, True)
      self.assertSequenceEqual(out2.shape, (bs, len_q, embed_dim))
  
file_path: scenic/scenic/projects/baselines/deformable_detr/tests/test_attention.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.deformable_detr.attention import bilinear_interpolate
  from scenic.projects.baselines.deformable_detr.attention import MultiScaleDeformableAttention
label_desc: Use image node_id 1 for calling the MultiScaleDeformableAttentionTest obj's underlying member method code with example usage: obj.test_ms_deformable_attn_output_shape(ref_dim, shapes, embed_dim, num_heads) without return types