code: def test_different_seq_lengths(self):
      query = jnp.zeros((2, 3))
      key = value = jnp.zeros((5, 3))
      mha = attention.MultiHeadAttention(
          key_size=7,
          num_heads=11,
          value_size=13,
          model_size=15,
          w_init_scale=1.0,
      )(query, key, value)
      self.assertEqual(mha.shape, (2, 15))
  
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 4 for calling the MultiHeadAttentionTest obj's underlying member method code with example usage: obj.test_different_seq_lengths() without return types