def _test_correctness_flex(B, H, S, D, mask_func, dtype, atol, rtol, device=infer_device()):
    """
    Test attention mechanisms with various implementations.

    Parameters:
        B (int): Batch size
        H (int): Number of attention heads
        S (int): Sequence length
        D (int): Hidden dimension per head
        mask_func: A function that generates custom attention mask
        dtype: Data type for computation
        atol (float): Absolute tolerance for comparison
        rtol (float): Relative tolerance for comparison
    """
    torch.manual_seed(0)

    # Initialize input tensors, i.e. the tensors after q, k, and v projections of hidden states (attention head input)
    query_torch = torch.randn(B, H, S, D, device=device, dtype=dtype, requires_grad=True)
    key_torch = torch.randn(B, H, S, D, device=device, dtype=dtype, requires_grad=True)
    value_torch = torch.randn(B, H, S, D, device=device, dtype=dtype, requires_grad=True)

    query_flex = query_torch.clone().detach().requires_grad_(True)
    key_flex = key_torch.clone().detach().requires_grad_(True)
    value_flex = value_torch.clone().detach().requires_grad_(True)

    block_mask = create_block_mask(mask_func, B, H, S, S, device=device)  # Sparsity block mask
    mask = create_mask(mask_func, B, H, S, S, device=device)  # Regular mask

    # If you are using a causal mask with FA2, you can enable `is_causal`."
    # e.g.,
    # F.scaled_dot_product_attention(query, key, value, is_causal=is_causal)

    torch_out = F.scaled_dot_product_attention(query_torch, key_torch, value_torch, attn_mask=mask)

    flex_out = flex_attention(query_flex, key_flex, value_flex, block_mask=block_mask)

    # Check forward pass
    assert_verbose_allclose(flex_out, torch_out, atol=atol, rtol=rtol)

    grad_out = torch.ones_like(torch_out)
    torch_out.backward(grad_out)
    flex_out.backward(grad_out)

    # Check gradients
    assert_verbose_allclose(query_flex.grad, query_torch.grad, atol=atol, rtol=rtol)
    assert_verbose_allclose(key_flex.grad, key_torch.grad, atol=atol, rtol=rtol)
    assert_verbose_allclose(value_flex.grad, value_torch.grad, atol=atol, rtol=rtol)

    Test attention mechanisms with various implementations.

    Parameters:
        B (int): Batch size
        H (int): Number of attention heads
        S (int): Sequence length
        D (int): Hidden dimension per head
        mask_func: A function that generates custom attention mask
        dtype: Data type for computation
        atol (float): Absolute tolerance for comparison
        rtol (float): Relative tolerance for comparison
    
test/transformers/test_flex_attention.py