code: def test_core_attention_mask_types():
      b, s, d = 4, 90, 16
      prob = 0.8  # make sure that we trigger the sparse kernels
  
      a = torch.rand(b, s, d)
      mask = torch.rand(b, s, s) > prob
  
      # mask of bools
      r_dense_bool = scaled_dot_product_attention(a, a, a, mask)
      r_sparse_bool = scaled_dot_product_attention(
          a, a, a, mask.to_sparse()
      )
      assert torch.allclose(r_dense_bool, r_sparse_bool)
  
      # Test additive mask. Mask of 0's and -infs.
      float_mask_add = torch.zeros_like(mask, dtype=torch.float)
      float_mask_add = float_mask_add.masked_fill(mask, float("-inf"))
  
      r_dense_add = scaled_dot_product_attention(
          a, a, a, float_mask_add
      )
      r_sparse_add = scaled_dot_product_attention(
          a, a, a, float_mask_add.to_sparse()
      )
  
      # Now properly handled
      assert torch.allclose(r_dense_add, r_sparse_add)
  
      # Test additive mask with mismatched batch dim
      d = b // 2
      mask = torch.rand(d, s, s) > prob
      float_mask_add = torch.zeros_like(mask, dtype=torch.float)
      float_mask_add = float_mask_add.masked_fill(mask, float("-inf"))
  
      # Make sure masking doesn't return errors
      r_dense_add = scaled_dot_product_attention(
          a, a, a, float_mask_add
      )
  
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 3 for calling a global function with example usage: test_core_attention_mask_types() without return types