code: def mha_reference(
      q,
      k,
      v,
      segment_ids: jnp.ndarray | None,
      sm_scale=1.0,
      causal: bool = False,
  ):
      q_seq_len = q.shape[1]
      kv_seq_len = k.shape[1]
      logits = jnp.einsum("bqhc,bkhc->bhqk", q, k).astype(jnp.float32)
      mask = None
      if segment_ids is not None:
          mask = jnp.expand_dims(
              segment_mask(segment_ids, segment_ids), 1
          )
          mask = jnp.broadcast_to(mask, logits.shape)
      if causal:
          causal_mask = jnp.tril(
              jnp.ones((1, 1, q_seq_len, kv_seq_len), dtype=bool)
          )
          causal_mask = jnp.broadcast_to(causal_mask, logits.shape)
          mask = (
              causal_mask
              if mask is None
              else jnp.logical_and(mask, causal_mask)
          )
      logits = (
          logits
          if mask is None
          else jnp.where(mask, logits, float("-inf"))
      )
      weights = jax.nn.softmax(logits * sm_scale).astype(q.dtype)
      return jnp.einsum("bhqk,bkhc->bqhc", weights, v)
  
file_path: jax/jax/experimental/pallas/ops/attention.py
imports: from __future__ import annotations
  import functools
  from typing import Any, Optional
  import jax
  from jax import lax
  from jax.experimental import pallas
  import jax.numpy
  import numpy
label_desc: Use image node_id 9 for calling a global function with example usage: mha_reference(q, k, v, segment_ids, sm_scale, causal) and returns: jnp