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: from __future__ import annotations
  import functools
  from typing import Any, Callable, Optional, Tuple, overload
  import jax
  import jax.numpy
  from jax import lax, random
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import rnglib
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import initializers
  from flax.experimental.nnx.nnx.nn.dtypes import promote_dtype
  from flax.experimental.nnx.nnx.nn.linear import DotGeneralT, LinearGeneral, PrecisionLike, default_kernel_init
  from flax.experimental.nnx.nnx.nn.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