code: def _check_head_dim_size_flash_nested(
      params: SDPAParams, debug=False
  ) -> bool:
      max_size = 256
      query_size_last = params.query.size(-1)
      key_size_last = params.key.size(-1)
      value_size_last = params.value.size(-1)
      same_head_dim_size = (
          query_size_last == key_size_last
          and query_size_last == value_size_last
      )
      if not (
          same_head_dim_size
          and (query_size_last % 8 == 0)
          and (query_size_last <= max_size)
      ):
          if debug:
              log.warning(
                  "For NestedTensor inputs, Flash attention requires q,k,v to have the same "
                  "last dimension and to be a multiple of 8 and less than or equal to 256. "
                  "Got Query.size(-1): %d, Key.size(-1): %d, Value.size(-1): %d instead.",
                  query_size_last,
                  key_size_last,
                  value_size_last,
              )
          return False
      return True
  
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 3 for calling a global function with example usage: _check_head_dim_size_flash_nested(params, debug) and returns: True, False