code: def create_tensors(
      shape, dtype, requires_grad=False, packed=True, multiquery=False
  ):
      stacked_shape = list(shape)  # B, M, H, K
      stacked_dim = 2 if packed else 0
      stacked_shape.insert(stacked_dim, 3)
      qkv = torch.rand(
          stacked_shape,
          device=device,
          dtype=dtype,
          requires_grad=requires_grad,
      )
      q = torch.rand(
          shape, device=device, dtype=dtype, requires_grad=requires_grad
      )
      shape_kv = (
          shape[0],
          shape[1],
          1 if multiquery else shape[2],
          shape[3],
      )
      k = torch.rand(
          shape_kv,
          device=device,
          dtype=dtype,
          requires_grad=requires_grad,
      ).expand(shape)
      v = torch.rand(
          shape_kv,
          device=device,
          dtype=dtype,
          requires_grad=requires_grad,
      ).expand(shape)
      return qkv, q, k, v
  
file_path: xformers/xformers/benchmarks/benchmark_mem_eff_attention.py
imports: import itertools
  import random
  from functools import partial
  import torch
  from torch.utils import benchmark
  from utils import benchmark_main_helper
  import xformers.ops
  import xformers.ops.fmha
  from xformers.attn_bias_utils import create_attn_bias
label_desc: Use image node_id 4 for calling a global function with example usage: create_tensors(shape, dtype, requires_grad, packed, multiquery) and returns: qkv,  q,  k,  v