class_docstr: Tests for modules in model_utils.py.
class_docstr_tok: ["Tests","for","modules","in","model_utils.py","."]
docstr: Tests PerformerEncoderAttention with dropout.
docstr_tok: ["Tests","PerformerEncoderAttention","with","dropout","."]
code: def test_performer_encoder_self_attention_w_dropout(
      self, attention_fn_cls
  ):
      """Tests PerformerEncoderAttention with dropout."""
      rng = random.PRNGKey(0)
      rng, dropout_rng = random.split(rng)
      x = jnp.ones((4, 16, 32))
      n_heads = 2
      encoder_self_attention_def = functools.partial(
          model_utils.PerformerEncoderAttention,
          num_heads=n_heads,
          attention_fn_cls=attention_fn_cls,
      )
      encoder_vars = encoder_self_attention_def().init(
          rng, x, x, deterministic=True
      )
      y = encoder_self_attention_def().apply(
          encoder_vars,
          x,
          x,
          deterministic=False,
          rngs={"dropout": dropout_rng},
      )
      # Test outputs shape.
      self.assertEqual(y.shape, x.shape)
  
file_path: scenic/scenic/projects/fast_vit/tests/test_model_utils.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.projects.fast_vit import model_utils
label_desc: Use image node_id 4 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_performer_encoder_self_attention_w_dropout(attention_fn_cls) without return types