code: def test_torch_script_ability(
      attention_name: str,
      heads: int,
      attn_dropout: float,
  ):
      if attention_name in {"favor", "global", "local", "random"}:
          # pyre-fixme[29]: The library function `pytest.skip` is not supported by Pyre.
          pytest.skip(
              f"{attention_name} does not support scripting yet."
          )
  
      device = torch.device("cpu")
  
      multi_head = _get_multihead(
          attention_name, attn_dropout, 0.0, False, heads, device
      )
  
      if (
          int(math.sqrt(SEQ)) ** 2 != SEQ
          and multi_head.attention.requires_squared_context
      ):
          pytest.skip(
              f"{attention_name} requires squared sequence lengths"
          )
  
      # input for tracing the function
      q = torch.rand((BATCH, SEQ, MODEL), device=device)
      k = torch.rand((BATCH, SEQ, MODEL), device=device)
      v = torch.rand((BATCH, SEQ, MODEL), device=device)
  
      # to make sure dropout behaves deterministically
      torch.random.manual_seed(42)
      # tracing the attention module
      traced_multi_head = torch.jit.trace(multi_head, (q, k, v))
  
      # create new random inputs for testing the eager model and traced model
      q = torch.rand((BATCH, SEQ, MODEL), device=device)
      k = torch.rand((BATCH, SEQ, MODEL), device=device)
      v = torch.rand((BATCH, SEQ, MODEL), device=device)
  
      # to make sure dropout behaves deterministically need to set the seed again
      torch.random.manual_seed(42)
      res = multi_head(query=q, key=k, value=v)
  
      # to make sure dropout behaves deterministically need to set the seed again
      torch.random.manual_seed(42)
      res_traced = traced_multi_head(query=q, key=k, value=v)
  
      assert torch.allclose(res, res_traced)
  
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 9 for calling a global function with example usage: test_torch_script_ability(attention_name, heads, attn_dropout) without return types