code: def _query_chunk_attention(
      query: Tensor,
      key: Tensor,
      value: Tensor,
      summarize_chunk: SummarizeChunk,
      kv_chunk_size: int,
  ) -> Tensor:
      batch_x_heads, k_tokens, k_channels_per_head = key.shape
      _, _, v_channels_per_head = value.shape
  
      def chunk_scanner(chunk_idx: int) -> AttnChunk:
          key_chunk = narrow_trunc(key, 1, chunk_idx, kv_chunk_size)
          value_chunk = narrow_trunc(value, 1, chunk_idx, kv_chunk_size)
          return summarize_chunk(query, key_chunk, value_chunk)
  
      chunks: list[AttnChunk] = [
          chunk_scanner(chunk)
          for chunk in torch.arange(0, k_tokens, kv_chunk_size)
      ]
      acc_chunk = AttnChunk(*map(torch.stack, zip(*chunks)))
      chunk_values, chunk_weights, chunk_max = acc_chunk
  
      global_max, _ = torch.max(chunk_max, 0, keepdim=True)
      max_diffs = torch.exp(chunk_max - global_max)
      chunk_values *= torch.unsqueeze(max_diffs, -1)
      chunk_weights *= max_diffs
  
      all_values = chunk_values.sum(dim=0)
      all_weights = torch.unsqueeze(chunk_weights, -1).sum(dim=0)
      return all_values / all_weights
  
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 3 for calling a global function with example usage: _query_chunk_attention(query, key, value, summarize_chunk, kv_chunk_size) and returns: unknown, summarize_chunk