code: def test_multihead_self_attention_w_dropout_disabled(self):
      rng = random.key(0)
      x = jnp.ones((4, 2, 3, 5))
      sa_module0 = nn.MultiHeadDotProductAttention(
          num_heads=8,
          qkv_features=16,
          kernel_init=initializers.ones,
          bias_init=initializers.zeros,
          dropout_rate=0.0,
          deterministic=True,
      )
      rng1, rng2, rng3, rng4 = random.split(rng, 4)
      rngs1 = {"params": rng1, "dropout": rng2}
      rngs2 = {"params": rng3, "dropout": rng4}
      y1, vs = sa_module0.init_with_output(rngs1, x)
      y2, _ = sa_module0.init_with_output(rngs2, x)
      np.testing.assert_allclose(y1, y2)
      y3 = sa_module0.apply(vs, x, rngs=rngs1)
      y4 = sa_module0.apply(vs, x, rngs=rngs2)
      np.testing.assert_allclose(y3, y4)
      sa_module1 = nn.MultiHeadDotProductAttention(
          num_heads=8,
          qkv_features=16,
          kernel_init=initializers.ones,
          bias_init=initializers.zeros,
          dropout_rate=0.0,
      )
      y5 = sa_module1.apply(vs, x, deterministic=True, rngs=rngs1)
      y6 = sa_module1.apply(vs, x, deterministic=True, rngs=rngs2)
      np.testing.assert_allclose(y5, y6)
      sa_module2 = nn.MultiHeadDotProductAttention(
          num_heads=8,
          qkv_features=16,
          kernel_init=initializers.ones,
          bias_init=initializers.zeros,
          dropout_rate=0.5,
      )
      y7 = sa_module2.apply(vs, x, deterministic=True, rngs=rngs1)
      y8 = sa_module2.apply(vs, x, deterministic=True, rngs=rngs2)
      np.testing.assert_allclose(y7, y8)
  
imports: import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest, parameterized
  from jax import lax, random
  from jax.nn import initializers
  from flax import errors, jax_utils
  from flax import linen
  from flax.core import pop
label_desc: Use image node_id 6 for calling the AttentionTest obj's underlying member method code with example usage: obj.test_multihead_self_attention_w_dropout_disabled() without return types