code: def _softmax_dispatch(
      x: torch.Tensor,
      log: bool,
      mask: Optional[torch.Tensor],
      causal: bool = False,
  ) -> torch.Tensor:
      # Triton is used if
      # - CUDA
      # - there's enough data to make it faster than pytorch. This could change over time, Triton is improving
      # - there was no previous failure
  
      global _triton_registered_warnings
  
      try:
          if (
              torch.cuda.is_available()
              and x.is_cuda
              and not _triton_registered_warnings
          ):
              return _softmax_triton.apply(x, mask, log, causal)
      except RuntimeError as e:
          # Catch cases where the current GPU does not have enough registers to hold a full tensor line
          # fallback to PyTorch's implementation, which streams the tensor in and out
          _triton_registered_warnings = True
          logger.warning(
              "Triton softmax kernel register spillover or invalid image caught."
              "Deactivating this kernel, please file an issue int the xFormers repository"
          )
          logger.warning(e)
  
      if mask is not None:
          x = x + mask
  
      if causal:
          x = x + torch.triu(
              torch.full_like(x, float("-inf")), diagonal=1
          )
  
      if log:
          return torch.log_softmax(x, dim=-1)
      else:
          return torch.softmax(x, dim=-1)
  
file_path: xformers/xformers/triton/softmax.py
imports: import logging
  from typing import Optional
  import torch
  import triton
  from torch.cuda.amp import custom_bwd, custom_fwd
  from xformers.triton.k_softmax import _softmax, _softmax_backward
label_desc: Use image node_id 3 for calling a global function with example usage: _softmax_dispatch(x, log, mask, causal) and returns: torch, torch, _softmax_triton