code: def test_shapes_single(
      self, seq_len, embed_size, d_key, num_heads, d_value, d_out
  ):
      query = key = value = jnp.zeros((seq_len, embed_size))
      mha = attention.MultiHeadAttention(
          key_size=d_key,
          num_heads=num_heads,
          value_size=d_value,
          model_size=d_out,
          w_init_scale=1.0,
      )(query, key, value)
      self.assertEqual(mha.shape, (seq_len, d_out))
  
file_path: dm-haiku/haiku/_src/attention_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import attention
  from haiku._src import initializers
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the MultiHeadAttentionTest obj's underlying member method code with example usage: obj.test_shapes_single(seq_len, embed_size, d_key, num_heads, d_value, d_out) without return types