code: def _get_attention_scores_no_kv_chunking(
      query: Tensor,
      key: Tensor,
      value: Tensor,
      scale: float,
  ) -> Tensor:
      attn_scores = torch.baddbmm(
          torch.zeros(1, 1, 1, device=query.device, dtype=query.dtype),
          query,
          key.transpose(1, 2),
          alpha=scale,
          beta=0,
      )
      attn_probs = attn_scores.softmax(dim=-1)
      del attn_scores
      hidden_states_slice = (
          torch.bmm(attn_probs, value)
          if query.device.type == "mps"
          else torch.bmm(attn_probs, value.to(attn_probs.dtype)).to(
              value.dtype
          )
      )
      return hidden_states_slice
  
file_path: stable-diffusion-webui/modules/sub_quadratic_attention.py
imports: from functools import partial
  import torch
  from torch import Tensor
  from torch.utils.checkpoint import checkpoint
  import math
  from typing import Optional, NamedTuple
label_desc: Use image node_id 4 for calling a global function with example usage: _get_attention_scores_no_kv_chunking(query, key, value, scale) and returns: hidden_states_slice