code: def test_dot_product_attention(getkey):
      q = jnp.array([[0.0, 2**0.5]])
      k = jnp.array([[1.0, 0.0], [0.0, 1.0]])
      v = jnp.array([[1.0], [0.0]])
      attn = eqx.nn._attention.dot_product_attention(q, k, v)
      assert attn.shape == (1, 1)
      assert jnp.allclose(attn, jnp.array([[1 / (1 + jnp.e)]]))
      mask = jnp.array([[True, False]])
      attn = eqx.nn._attention.dot_product_attention(q, k, v, mask)
      assert attn == jnp.array([[1.0]])
  
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 22 for calling a global function with example usage: test_dot_product_attention(getkey) without return types