code: def test_multihead_attention_inference(getkey):
      attn = eqx.nn.MultiheadAttention(
          num_heads=2, query_size=10, dropout_p=0.5, key=getkey()
      )
      x = jrandom.normal(getkey(), (3, 10))
      z1 = attn(x, x, x, key=getkey(), inference=True)
      z2 = attn(x, x, x, inference=True)
      assert jnp.all(z1 == z2)
  
      z1 = attn(x, x, x, key=getkey())
      z2 = attn(x, x, x, key=getkey())
      assert jnp.any(z1 != z2)
  
      attn2 = eqx.nn.MultiheadAttention(
          num_heads=2,
          query_size=10,
          dropout_p=0.5,
          inference=True,
          key=getkey(),
      )
      z1 = attn2(x, x, x, key=getkey())
      z2 = attn2(x, x, x, key=getkey())
      assert jnp.all(z1 == z2)
  
imports: import warnings
  from typing import Union
  import equinox
  import jax
  import jax.nn
  import jax.numpy
  import jax.random
  import pytest
label_desc: Use image node_id 24 for calling a global function with example usage: test_multihead_attention_inference(getkey) without return types