code: def test_different_kq_dimensions(
      attention_name: str,
      heads: int,
      device: torch.device,
  ):
      multi_head = _get_multihead(
          attention_name, 0.0, 0.0, False, heads, device
      )
  
      if multi_head.attention.requires_same_k_q_dimensions:
          # pyre-fixme[29]: The library function `pytest.skip` is not supported by Pyre.
          pytest.skip(
              f"{attention_name} does not support different k, q dimensions yet."
          )
  
      seq_q = SEQ // 2
      q = torch.rand((BATCH, seq_q, MODEL), device=device)
      k = torch.rand((BATCH, SEQ, MODEL), device=device)
      v = torch.rand((BATCH, SEQ, MODEL), device=device)
  
      res = multi_head(query=q, key=k, value=v)
      assert res.shape == torch.Size([BATCH, seq_q, MODEL])
  
file_path: xformers/tests/test_attentions.py
imports: import math
  from typing import Tuple
  import pytest
  import torch
  from xformers.components import InputProjection, InputProjectionConfig, MultiHeadDispatch
  from xformers.components.attention import _DENSITY_THRESHOLD, ATTENTION_REGISTRY, build_attention
label_desc: Use image node_id 6 for calling a global function with example usage: test_different_kq_dimensions(attention_name, heads, device) without return types