code: def make_attention_mask(
      query_input: Array,
      key_input: Array,
      pairwise_fn: Callable[..., Any] = jnp.multiply,
      extra_batch_dims: int = 0,
      dtype: Dtype = jnp.float32,
  ):
      """Mask-making helper for attention weights.
  
      In case of 1d inputs (i.e., ``[batch..., len_q]``, ``[batch..., len_kv]``, the
      attention weights will be ``[batch..., heads, len_q, len_kv]`` and this
      function will produce ``[batch..., 1, len_q, len_kv]``.
  
      Args:
        query_input: a batched, flat input of query_length size
        key_input: a batched, flat input of key_length size
        pairwise_fn: broadcasting elementwise comparison function
        extra_batch_dims: number of extra batch dims to add singleton axes for, none
          by default
        dtype: mask return dtype
  
      Returns:
        A ``[batch..., 1, len_q, len_kv]`` shaped mask for 1d attention.
      """
      mask = pairwise_fn(
          jnp.expand_dims(query_input, axis=-1),
          jnp.expand_dims(key_input, axis=-2),
      )
      mask = jnp.expand_dims(mask, axis=-3)
      mask = jnp.expand_dims(mask, axis=tuple(range(extra_batch_dims)))
      return mask.astype(dtype)
  
imports: import functools
  import warnings
  from typing import Any, Callable, Optional, Tuple, Union, overload
  import jax
  import jax.numpy
  from jax import lax, random
  from flax.linen import initializers
  from flax.linen.dtypes import promote_dtype
  from flax.linen.linear import DenseGeneral, DotGeneralT, PrecisionLike, default_kernel_init
  from flax.linen.module import Module, compact, merge_param
  from flax.linen.normalization import LayerNorm
label_desc: Use image node_id 3 for calling a global function with example usage: make_attention_mask(query_input, key_input, pairwise_fn, extra_batch_dims, dtype) and returns: mask