code: def dot_product_attention_weights(
      query: Float[Array, "q_seq qk_size"],
      key: Float[Array, "kv_seq qk_size"],
      mask: Optional[Bool[Array, "q_seq kv_seq"]] = None,
  ) -> Float[Array, "q_seq kv_seq"]:
      query = query / math.sqrt(query.shape[-1])
      logits = jnp.einsum("sd,Sd->sS", query, key)
      if mask is not None:
          if mask.shape != logits.shape:
              raise ValueError(
                  f"mask must have shape (query_seq_length, "
                  f"kv_seq_length)=({query.shape[0]}, "
                  f"{key.shape[0]}). Got {mask.shape}."
              )
          logits = jnp.where(mask, logits, jnp.finfo(logits.dtype).min)
          logits = cast(Array, logits)
  
      with jax.numpy_dtype_promotion("standard"):
          dtype = jnp.result_type(logits.dtype, jnp.float32)
      weights = jax.nn.softmax(logits.astype(dtype)).astype(
          logits.dtype
      )
      return weights
  
file_path: equinox/equinox/nn/_attention.py
imports: import functools
  import math
  import warnings
  from functools import partial
  from typing import cast, Optional, Union
  import jax
  import jax.numpy
  import jax.random
  from jaxtyping import Array, Bool, Float, PRNGKeyArray
  from .._module import field, Module
  from ._dropout import Dropout
  from ._linear import Linear
label_desc: Use image node_id 1 for calling a global function with example usage: dot_product_attention_weights(query, key, mask) and returns: weights