code: def construct_local_mask(
      seqlen_q,
      seqlen_k,
      window_size=(-1, -1),  # -1 means infinite window size
      query_padding_mask=None,
      key_padding_mask=None,
      device=None,
  ):
      row_idx = rearrange(
          torch.arange(seqlen_q, device=device, dtype=torch.long),
          "s -> s 1",
      )
      col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
      sk = (
          seqlen_k
          if key_padding_mask is None
          else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
      )
      sq = (
          seqlen_q
          if query_padding_mask is None
          else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
      )
      if window_size[0] < 0:
          return col_idx > row_idx + sk - sq + window_size[1]
      else:
          sk = (
              torch.full_like(col_idx, seqlen_k)
              if key_padding_mask is None
              else sk
          )
          return torch.logical_or(
              col_idx
              > torch.minimum(row_idx + sk - sq + window_size[1], sk),
              col_idx < row_idx + sk - sq - window_size[0],
          )
  
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from flash_attn import flash_attn_func, flash_attn_kvpacked_func, flash_attn_qkvpacked_func, flash_attn_varlen_func, flash_attn_varlen_kvpacked_func, flash_attn_varlen_qkvpacked_func, flash_attn_with_kvcache
  from flash_attn.bert_padding import pad_input, unpad_input
  from flash_attn.flash_attn_interface import _get_block_size_n
  from flash_attn.layers.rotary import apply_rotary_emb
label_desc: Use image node_id 4 for calling a global function with example usage: construct_local_mask(seqlen_q, seqlen_k, window_size, query_padding_mask, key_padding_mask, device) and returns: unknown, torch