code: def make_attention_mask(
      query_input: tp.Any,
      key_input: tp.Any,
      pairwise_fn: tp.Callable = jnp.multiply,
      dtype: tp.Any = jnp.float32,
  ):
      mask = pairwise_fn(
          jnp.expand_dims(query_input, axis=-1),
          jnp.expand_dims(key_input, axis=-2),
      )
      return jnp.expand_dims(mask, axis=-3).astype(dtype)
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling a global function with example usage: make_attention_mask(query_input, key_input, pairwise_fn, dtype) and returns: jnp