code: def _check_batch_size_nested(params: SDPAParams, debug=False) -> bool:
      # This is expected to be called after check_tensor_shapes ensuring that the
      # size() calls won't error since the inputs are all 4 dimensional
      q_batch_size = params.query.size(0)
      k_batch_size = params.key.size(0)
      v_batch_size = params.value.size(0)
  
      # num_heads logic for nested input is checked in
      # check_for_seq_len_0_nested_tensor as there is handling there to make sure
      # num_heads is not ragged
      return (
          q_batch_size == k_batch_size and q_batch_size == v_batch_size
      )
  
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 2 for calling a global function with example usage: _check_batch_size_nested(params, debug) and returns: unknown