code: def _mha_reference(
      q,
      k,
      v,
      ab,
      segment_ids: SegmentIds | None,
      causal: bool,
      mask_value: float,
      sm_scale: float,
      save_residuals: bool,
  ):
      return mha_reference_no_custom_vjp(
          q,
          k,
          v,
          ab,
          segment_ids,
          causal=causal,
          mask_value=mask_value,
          sm_scale=sm_scale,
          save_residuals=save_residuals,
      )
  
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 16 for calling a global function with example usage: _mha_reference(q, k, v, ab, segment_ids, causal, mask_value, sm_scale, save_residuals) and returns: mha_reference_no_custom_vjp