code: def test_switch_blocksparse_dropout(device, training, drop_prob):
      b, s, d = 8, 128, 32
  
      a = torch.rand(b, s, d, device=device)
  
      m = AttentionMask.make_causal(s, s, device)
      dropout = nn.Dropout(drop_prob)
      dropout.train(training).cuda()
  
      with torch.cuda.amp.autocast():
          r = scaled_dot_product_attention(a, a, a, m)
          r_drop = scaled_dot_product_attention(a, a, a, m, dropout)
  
      # Check for dropout when applicable
      if dropout.p and dropout.training:
          assert (r_drop != r).any()
      else:
          assert torch.allclose(r, r_drop)
  
file_path: xformers/tests/test_core_attention.py
imports: import functools
  import pytest
  import torch
  from torch import nn
  from xformers import _is_triton_available
  from xformers.components.attention._sputnik_sparse import SparseCS
  from xformers.components.attention.attention_mask import AttentionMask
  from xformers.components.attention.core import scaled_dot_product_attention
label_desc: Use image node_id 10 for calling a global function with example usage: test_switch_blocksparse_dropout(device, training, drop_prob) without return types