code: def test_with_bias(self, with_bias, expected_params):
      def f(key, query, value):
          w_init = initializers.Constant(3)
          mha1 = attention.MultiHeadAttention(
              2, 3, w_init=w_init, with_bias=with_bias
          )
          return mha1(key, query, value)
  
      rng = jax.random.PRNGKey(42)
      init, _ = transform.transform(f)
      query = key = jnp.zeros((5, 3))
      value = jnp.zeros((5, 10))
      params = init(rng, key, query, value)
      for module_params in params.values():
          self.assertLen(module_params, expected_params)
  
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 9 for calling the MultiHeadAttentionTest obj's underlying member method code with example usage: obj.test_with_bias(with_bias, expected_params) and returns: mha1