code: def test_w_init(self):
      with self.assertRaisesRegex(
          ValueError, "provide a weight initializer"
      ):
          attention.MultiHeadAttention(2, 3)
      with self.assertRaisesRegex(ValueError, "provide only `w_init`"):
          attention.MultiHeadAttention(
              2, 3, w_init_scale=5, w_init=initializers.Constant(0)
          )
  
      w_init = initializers.Constant(3)
      mha1 = attention.MultiHeadAttention(2, 3, w_init=w_init)
      self.assertIs(mha1.w_init, w_init)
  
      mha2 = attention.MultiHeadAttention(2, 3, w_init_scale=5)
      self.assertIsInstance(mha2.w_init, initializers.VarianceScaling)
  
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 7 for calling the MultiHeadAttentionTest obj's underlying member method code with example usage: obj.test_w_init() without return types