code: def _mha_forward(
      q,
      k,
      v,
      segment_ids: jax.Array | None,
      sm_scale: float,
      causal: bool,
      block_q: int,
      block_k: int,
      backward_pass_impl: str,
      num_warps: int | None,
      num_stages: int,
      grid: Any,
      interpret: bool,
      debug: bool,
  ):
      del backward_pass_impl
      batch_size, seq_len, num_heads, head_dim = q.shape
      block_q = min(block_q, seq_len)
      block_k = min(block_k, seq_len)
      # Heuristics.
      grid_ = grid
      if grid_ is None:
          grid_ = (pl.cdiv(seq_len, block_q), batch_size, num_heads)
  
      num_warps_ = num_warps
      if num_warps_ is None:
          num_warps_ = 4 if head_dim <= 64 else 8
      kernel = functools.partial(
          mha_forward_kernel,
          num_heads=num_heads,
          sm_scale=sm_scale,
          causal=causal,
          block_q=block_q,
          block_k=block_k,
          block_d=head_dim,
      )
      out_shape = [
          jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype),  # out
          jax.ShapeDtypeStruct(
              shape=(batch_size, num_heads, seq_len),  # l
              dtype=jnp.float32,
          ),
          jax.ShapeDtypeStruct(
              shape=(batch_size, num_heads, seq_len),  # m
              dtype=jnp.float32,
          ),
      ]
      in_specs = [
          pl.BlockSpec(
              lambda _, j, k: (j, 0, k, 0),
              (None, seq_len, None, head_dim),
          ),
          pl.BlockSpec(
              lambda _, j, k: (j, 0, k, 0),
              (None, seq_len, None, head_dim),
          ),
          pl.BlockSpec(
              lambda _, j, k: (j, 0, k, 0),
              (None, seq_len, None, head_dim),
          ),
      ]
      in_specs.append(
          None  # type: ignore[arg-type]
          if segment_ids is None
          else pl.BlockSpec(lambda _, j, k: (j, 0), (None, seq_len))
      )
      out, l, m = pl.pallas_call(
          kernel,
          grid=grid_,
          in_specs=in_specs,
          out_specs=[
              pl.BlockSpec(
                  lambda _, j, k: (j, 0, k, 0),
                  (None, seq_len, None, head_dim),
              ),
              pl.BlockSpec(
                  lambda _, j, k: (j, k, 0), (None, None, seq_len)
              ),
              pl.BlockSpec(
                  lambda _, j, k: (j, k, 0), (None, None, seq_len)
              ),
          ],
          num_warps=num_warps_,
          num_stages=num_stages,
          out_shape=out_shape,
          debug=debug,
          interpret=interpret,
          name="mha_forward",
      )(q, k, v, segment_ids)
      return out, (q, k, v, segment_ids, out, l, m)
  
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 4 for calling a global function with example usage: _mha_forward(q, k, v, segment_ids, sm_scale, causal, block_q, block_k, backward_pass_impl, num_warps, num_stages, grid, interpret, debug) and returns: out,  q,  k,  v,  segment_ids,  out,  l,  m