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(self):
      """Tests MultiHeadAttention."""
      rng = random.PRNGKey(0)
      x = jnp.ones((4, 16, 32))
      n_heads = 2
      layer = attention_layers.MultiHeadAttention(num_heads=n_heads)
      variables = layer.init(rng, x, x, deterministic=True)
      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 2 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_multihead_attention() without return types