code: def _summarize_chunk(
      query: Tensor,
      key: Tensor,
      value: Tensor,
      scale: float,
  ) -> AttnChunk:
      attn_weights = torch.baddbmm(
          torch.zeros(1, 1, 1, device=query.device, dtype=query.dtype),
          query,
          key.transpose(1, 2),
          alpha=scale,
          beta=0,
      )
      max_score, _ = torch.max(attn_weights, -1, keepdim=True)
      max_score = max_score.detach()
      exp_weights = torch.exp(attn_weights - max_score)
      exp_values = (
          torch.bmm(exp_weights, value)
          if query.device.type == "mps"
          else torch.bmm(exp_weights, value.to(exp_weights.dtype)).to(
              value.dtype
          )
      )
      max_score = max_score.squeeze(-1)
      return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)
  
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 2 for calling a global function with example usage: _summarize_chunk(query, key, value, scale) and returns: AttnChunk