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