code: def test_softmax_parity_fallback(
      log, masking, causal, contiguous, device
  ):
      """Check that the fallback paths are correct"""
      torch.random.manual_seed(0)
  
      shape = (16, 16)
  
      # Check the result of a FW pass
      X = torch.normal(
          0, 1, size=shape, device=device, requires_grad=False
      )
  
      if not contiguous:
          # Make sure that the buffer is not contiguous
          X = X.transpose(-2, -1).contiguous().transpose(-2, -1)
  
      X_ = X.clone()
      X.requires_grad = True
      X_.requires_grad = True
  
      seq = shape[1]
      mask = torch.zeros((seq, seq), device=device)
      if masking:
          mask[torch.rand((seq, seq), device=device) > 0.8] = -float(
              "inf"
          )
  
      mask_causal = torch.zeros_like(mask)
      if causal:
          mask_causal[
              ~torch.tril(torch.ones_like(mask)).bool()
          ] = -float("inf")
  
      y_torch = (
          torch.log_softmax(X + mask + mask_causal, dim=-1)
          if log
          else torch.softmax(X + mask + mask_causal, dim=-1)
      )
      y_triton = (
          triton_log_softmax(X_, mask, causal)
          if log
          else triton_softmax(X_, mask, causal)
      )
  
      assert torch.allclose(y_torch, y_triton, equal_nan=True)
  
      # Check that BW also gives the same result
      loss_torch = torch.norm(y_torch.transpose(-2, -1) @ y_torch)
      loss_torch.backward()
  
      loss_triton = torch.norm(y_triton.transpose(-2, -1) @ y_triton)
      loss_triton.backward()
  
      assert torch.allclose(
          torch.norm(X.grad),
          torch.norm(X_.grad),
          equal_nan=True,
          atol=1e-5,
      ), f"{torch.norm(X.grad)}, {torch.norm(X_.grad)}"
  
file_path: xformers/tests/test_triton_softmax.py
imports: import logging
  import pytest
  import torch
  from torch.cuda.amp.autocast_mode import autocast
  import xformers
label_desc: Use image node_id 3 for calling a global function with example usage: test_softmax_parity_fallback(log, masking, causal, contiguous, device) without return types