code: def mha_reference_no_custom_vjp(
      q,
      k,
      v,
      ab: jax.Array | None = None,
      segment_ids: SegmentIds | None = None,
      *,
      causal: bool = False,
      mask_value: float = DEFAULT_MASK_VALUE,
      sm_scale: float = 1.0,
      save_residuals: bool = False,
  ):
      logits = jnp.einsum("bhqc,bhkc->bhqk", q, k)
      if ab is not None:
          logits += ab
      if sm_scale != 1.0:
          logits *= sm_scale
  
      mask = None
      if segment_ids is not None:
          mask = segment_ids.q[:, :, None] == segment_ids.kv[:, None, :]
          mask = mask[:, None, :, :]
  
      if causal:
          _, _, q_seq_len, _ = q.shape
          _, _, kv_seq_len, _ = k.shape
          mask_shape = (q_seq_len, kv_seq_len)
          row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0)
          col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
          causal_mask = (col_ids <= row_ids)[None, None, :, :]
          mask = (
              causal_mask
              if mask is None
              else jnp.logical_and(mask, causal_mask)
          )
  
      logits = (
          logits
          if mask is None
          else logits + jnp.where(mask, 0.0, mask_value)
      )
  
      m = logits.max(axis=-1)
      unnormalized = jnp.exp(logits - m[..., None])
      l = unnormalized.sum(axis=-1)
      weights = unnormalized / l[..., None]
      out = jnp.einsum("bhqk,bhkc->bhqc", weights, v)
      if save_residuals:
          return out, l, m
      return out
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  from typing import Any, NamedTuple
  import jax
  from jax import lax
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  import jax.numpy
label_desc: Use image node_id 14 for calling a global function with example usage: mha_reference_no_custom_vjp(q, k, v, ab, segment_ids) and returns: out, out,  l,  m