class_docstr: Tests for modules in attention_layers.py.
class_docstr_tok: ["Tests","for","modules","in","attention_layers.py","."]
docstr: Tests MultiHeadAttention with dropout.
docstr_tok: ["Tests","MultiHeadAttention","with","dropout","."]
code: def test_multihead_attention_w_dropout(self):
      """Tests MultiHeadAttention with dropout."""
      rng = random.PRNGKey(0)
      rng, dropout_rng = random.split(rng)
      x = jnp.ones((4, 16, 32))
      n_heads = 2
      layer = attention_layers.MultiHeadAttention(
          num_heads=n_heads, dropout_rate=0.1
      )
      variables = layer.init(rng, x, x, deterministic=True)
      y = layer.apply(
          variables,
          x,
          x,
          deterministic=False,
          rngs={"dropout": dropout_rng},
      )
      # 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 4 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_multihead_attention_w_dropout() without return types