code: def test_dropout(shape, amp, bias, p):
      """
      Check some basic dropout properties
      """
      torch.random.manual_seed(0)
      torch.cuda.manual_seed_all(0)
  
      x = torch.normal(
          0, 1, size=shape, device="cuda", requires_grad=True
      )
      b = (
          torch.normal(
              0, 1, size=(shape[-1],), device="cuda", requires_grad=True
          )
          if bias
          else None
      )
  
      with autocast(enabled=amp):
          tol = (
              1e-2 if amp else 1e-5
          )  # AMP rounding causes issues, 1e-5 is the default
  
          # Check that 0 means no dropout
          y = triton_dropout(x, p=0, bias=b)
          x_ref = (x + b if bias else x).to(y.dtype)
          assert torch.allclose(x_ref, y, rtol=tol), f"{x[x>y]}"
  
          # Check that 1 means drop all
          y = triton_dropout(x, p=1, bias=b)
          x_ref = (x + b if bias else x).to(y.dtype)
          assert torch.allclose(torch.zeros_like(y), y, rtol=tol)
  
          # Check that .99 means probably dropout
          y = triton_dropout(x, p=0.99, bias=b)
          x_ref = (x + b if bias else x).to(y.dtype)
          assert not torch.allclose(x_ref, y, rtol=tol)
  
          # Check that the drops are different for every row (could catch broken seeds per row)
          y = triton_dropout(x, p=0.5)
  
          y = y.flatten(0, 1) if y.ndim == 3 else y
          assert (
              not torch.sum(torch.eq(y[0, :] == 0.0, y[1, :] == 0.0))
              == y.shape[1]
          )
  
          # Check that the drops are different over time, for the same line
          y_a = triton_dropout(x, p=0.5)
          y_b = triton_dropout(x, p=0.5)
  
          y_a = y_a.flatten(0, 1) if y_a.ndim == 3 else y_a
          y_b = y_b.flatten(0, 1) if y_b.ndim == 3 else y_b
  
          assert (
              not torch.sum(
                  torch.eq(y_a[0, :] == 0.0, y_b[0, :] == 0.0)
              ).item()
              == y.shape[1]
          )
  
          # Check that the drop probability is about right
          y = triton_dropout(x, p=p)
          drop_p = (y.numel() - y.count_nonzero()) / y.numel()
          assert abs(drop_p - p) < 0.02
  
          # Check that the same seeds lead to the same dropout
          torch.manual_seed(0)
          torch.cuda.manual_seed(0)
          y_1 = triton_dropout(x, p=0.5)
  
          torch.manual_seed(0)
          torch.cuda.manual_seed(0)
          y_2 = triton_dropout(x, p=0.5)
  
          torch.testing.assert_close(y_1, y_2)
  
file_path: xformers/tests/test_triton_dropout.py
imports: import logging
  import pytest
  import torch
  from torch.cuda.amp.autocast_mode import autocast
  import xformers
  from xformers.components import Activation, build_activation
label_desc: Use image node_id 2 for calling a global function with example usage: test_dropout(shape, amp, bias, p) without return types