code: def test_dot_product_attention_weights(getkey):
      q = jnp.array([[0.0, 2**0.5]])
      k = jnp.array([[1.0, 0.0], [0.0, 1.0]])
      weights = eqx.nn._attention.dot_product_attention_weights(q, k)
      assert weights.shape == (1, 2)
      assert jnp.allclose(
          weights, jnp.array([[1, jnp.e]]) / (1 + jnp.e)
      )
      mask = jnp.array([[True, False]])
      weights = eqx.nn._attention.dot_product_attention_weights(
          q, k, mask
      )
      assert jnp.allclose(weights, jnp.array([[1.0, 0.0]]))
  
      q = jnp.array([[1.0]], dtype="float16")
      k = jnp.array([[9.0], [1.0]], dtype="float16")
      weights = eqx.nn._attention.dot_product_attention_weights(q, k)
      assert weights.dtype == q.dtype
      assert weights.max() < 1
  
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 21 for calling a global function with example usage: test_dot_product_attention_weights(getkey) without return types