code: def _mha_reference_fwd(
      q,
      k,
      v,
      ab,
      segment_ids: SegmentIds | None,
      causal: bool,
      mask_value: float,
      sm_scale: float,
      save_residuals: bool,
  ):
      if save_residuals:
          raise NotImplementedError
      res = _mha_reference(
          q,
          k,
          v,
          ab,
          segment_ids,
          causal=causal,
          mask_value=mask_value,
          sm_scale=sm_scale,
          save_residuals=True,
      )
      assert isinstance(res, tuple)
      out, l, m = res
      return out, (q, k, v, ab, segment_ids, out, l, m)
  
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 17 for calling a global function with example usage: _mha_reference_fwd(q, k, v, ab, segment_ids, causal, mask_value, sm_scale, save_residuals) and returns: out,  q,  k,  v,  ab,  segment_ids,  out,  l,  m