code: def _flash_attention_fwd(
      q,
      k,
      v,
      ab,
      segment_ids,
      save_residuals,
      causal,
      sm_scale,
      block_sizes,
      debug,
  ):
      if save_residuals:
          raise NotImplementedError("Higher-order AD not supported")
      o, l, m = _flash_attention(
          q,
          k,
          v,
          ab,
          segment_ids,
          True,
          causal,
          sm_scale,
          block_sizes,
          debug,
      )
      return o, (q, k, v, ab, segment_ids, o, 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 3 for calling a global function with example usage: _flash_attention_fwd(q, k, v, ab, segment_ids, save_residuals, causal, sm_scale, block_sizes, debug) and returns: o,  q,  k,  v,  ab,  segment_ids,  o,  l,  m