class_docstr: Tests for modules in attention_layers.py.
class_docstr_tok: ["Tests","for","modules","in","attention_layers.py","."]
docstr: Tests MultiHeadAttention.
docstr_tok: ["Tests","MultiHeadAttention","."]
code: def test_multihead_attention_hidden_size_not_divisible_by_heads(self):
      """Tests MultiHeadAttention."""
      rng = random.PRNGKey(0)
      x = jnp.ones((4, 16, 30))
      n_heads = 4
      layer = attention_layers.MultiHeadAttention(
          num_heads=n_heads,
          enforce_hidden_size_divisible_by_heads=False,
      )
      variables = layer.init(rng, x, x, deterministic=True)
      self.assertTupleEqual(
          variables["params"]["query"]["kernel"].shape, (30, 4, 7)
      )
      self.assertTupleEqual(
          variables["params"]["key"]["kernel"].shape, (30, 4, 7)
      )
      self.assertTupleEqual(
          variables["params"]["value"]["kernel"].shape, (30, 4, 7)
      )
      self.assertTupleEqual(
          variables["params"]["out"]["kernel"].shape, (4, 7, 30)
      )
      y = layer.apply(variables, x, x, deterministic=True)
      # Test outputs shape.
      self.assertEqual(y.shape, x.shape)
  
file_path: scenic/scenic/model_lib/layers/tests/test_attention_layers.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 3 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_multihead_attention_hidden_size_not_divisible_by_heads() without return types