code: def jagged_scaled_dot_product_attention(
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      attn_mask: Optional[torch.Tensor] = None,
      dropout_p=0.0,
      is_causal=False,
      scale=None,
  ):
      _validate_sdpa_input(
          query, key, value, attn_mask, dropout_p, is_causal, scale
      )
      # for mypy, ugh
      assert (
          isinstance(query, NestedTensor)
          and isinstance(key, NestedTensor)
          and isinstance(value, NestedTensor)
      )
  
      # Special path for non-ragged sequence length (e.g. for SAM where we have a ragged
      # second batch dim instead). For this case, we can just send the dense buffers through
      # vanilla SDPA.
      if (
          query.dim() > 3
          and key.dim() > 3
          and value.dim() > 3
          and query._ragged_idx == 1
      ):
          from torch.nested._internal.ops import extract_kwargs
  
          output = F.scaled_dot_product_attention(
              query._values,
              key._values,
              value._values,
              attn_mask=(
                  attn_mask._values
                  if isinstance(attn_mask, NestedTensor)
                  else attn_mask
              ),
              dropout_p=dropout_p,
              is_causal=is_causal,
              scale=scale,
          )
  
          return NestedTensor(output, **extract_kwargs(query))
  
      compute_logsumexp = (
          query.requires_grad
          or key.requires_grad
          or value.requires_grad
      )
  
      backend_choice = _select_sdp_backend(
          query, key, value, attn_mask, dropout_p, is_causal
      )
  
      if backend_choice == SDPBackend.FLASH_ATTENTION:
          og_size = query.size(-1)
          query_padded = _pad_last_dim(query, 8, False)
          key_padded = _pad_last_dim(key, 8, False)
          value_padded = _pad_last_dim(value, 8, False)
          # We need to calculate the scale based off the OG head dim size
          og_scale = _calculate_scale(query, scale)
          (
              query_buffer_reshaped,
              key_buffer_reshaped,
              value_buffer_reshaped,
              cumulative_sequence_length_q,
              cumulative_sequence_length_kv,
              max_seqlen_batch_q,
              max_seqlen_batch_kv,
              output_nt_info,
          ) = _sdpa_nested_preprocessing(
              query_padded, key_padded, value_padded
          )
  
          (
              attention,
              logsumexp,
              philox_seed,
              philox_offset,
              debug_attn_mask,
          ) = torch.ops.aten._flash_attention_forward(
              query_buffer_reshaped,
              key_buffer_reshaped,
              value_buffer_reshaped,
              cumulative_sequence_length_q,
              cumulative_sequence_length_kv,
              max_seqlen_batch_q,
              max_seqlen_batch_kv,
              dropout_p,
              is_causal,
              False,
              scale=og_scale,
          )
          # Reshape output to convert nnz to batch_size and seq_len
          attention = ViewNestedFromBuffer.apply(
              attention.squeeze(0), output_nt_info["offsets"]
          ).transpose(1, 2)
          return _post_process_flash_output(attention, og_size)
      elif backend_choice == SDPBackend.EFFICIENT_ATTENTION:
          (
              query_reshaped,
              key_reshaped,
              value_reshaped,
              cumulative_sequence_length_q,
              cumulative_sequence_length_kv,
              max_seqlen_batch_q,
              max_seqlen_batch_kv,
              output_nt_info,
          ) = _sdpa_nested_preprocessing(query, key, value)
          (
              attention,
              log_sumexp,
              seed,
              offset,
              max_seqlen_q,
              max_seqlen_batch_kv,
          ) = torch.ops.aten._efficient_attention_forward(
              query_reshaped.unsqueeze(0),
              key_reshaped.unsqueeze(0),
              value_reshaped.unsqueeze(0),
              None,
              cumulative_sequence_length_q,
              cumulative_sequence_length_kv,
              max_seqlen_batch_q,
              max_seqlen_batch_kv,
              dropout_p,
              int(is_causal),
              compute_logsumexp,
              scale=scale,
          )
  
          # Reshape output to convert nnz to batch_size and seq_len
          return ViewNestedFromBuffer.apply(
              attention.squeeze(0), output_nt_info["offsets"]
          ).transpose(1, 2)
      elif backend_choice == SDPBackend.MATH:
          # save the offsets and shape of the inputs, so we can reshape the final output
          # query @ key = attn: [B, D1, j0, D'] @ [B, D1, D' j1] = [B, D1, j0, j1]
          # attn @ value = out: [B, D1, j0, j1] @ [B, D1, j1, D2] = [B, D1, j0, D2]
          offsets = query.offsets()
          d1 = query._size[1]
          d2 = value._size[-1]
  
          # convert jagged layout Nested Tensor to strided layout Nested Tensor
          # which support the math implementation of SDPA
          def get_strided_layout_nested_tensor(jagged_layout_nt):
              lengths = (
                  jagged_layout_nt._offsets[1:]
                  - jagged_layout_nt._offsets[:-1]
              )
              transpose = torch.transpose(jagged_layout_nt, 1, 2)
              tensor_list = buffer_from_jagged(transpose).split(
                  list(lengths), dim=0
              )
              strided_nt = torch.nested.as_nested_tensor(
                  list(tensor_list)
              )
              strided_nt = strided_nt.transpose(1, 2).contiguous()
              return strided_nt
  
          query = get_strided_layout_nested_tensor(query)
          key = get_strided_layout_nested_tensor(key)
          value = get_strided_layout_nested_tensor(value)
  
          attn_out = torch._scaled_dot_product_attention_math(
              query,
              key,
              value,
              attn_mask,
              dropout_p,
              is_causal,
              scale=scale,
          )[0]
  
          # convert strided layout Nested Tensor back to jagged layout Nested Tensor
          attn_out = attn_out.transpose(1, 2).contiguous().values()
          attn_out = attn_out.view(-1, d1, d2)
          attn_out = ViewNestedFromBuffer.apply(attn_out, offsets)
          attn_out = attn_out.transpose(1, 2)
  
          return attn_out
      else:
          raise RuntimeError(
              "No viable backend for scaled_dot_product_attention was found."
          )
  
imports: import logging
  import math
  from typing import Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.backends.cuda import can_use_efficient_attention, can_use_flash_attention, flash_sdp_enabled, math_sdp_enabled, mem_efficient_sdp_enabled, SDPAParams, SDPBackend
  from .nested_tensor import buffer_from_jagged, NestedTensor, ViewNestedFromBuffer
label_desc: Use image node_id 18 for calling a global function with example usage: jagged_scaled_dot_product_attention(query, key, value, attn_mask, dropout_p, is_causal, scale) and returns: NestedTensor, _post_process_flash_output, ViewNestedFromBuffer, attn_out, strided_nt