code: def softmax(
      x: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
      causal: bool = False,
  ) -> torch.Tensor:
      r"""Applies the Softmax function to an 3-dimensional input Tensor
      rescaling them so that the elements of the n-dimensional output Tensor
      lie in the range [0,1] and sum to 1.
  
      Softmax is defined as:
  
      .. math::
          \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
  
      .. warning: softmax is computed on the last dimension of the input tensor.
  
  
      Args:
          x: input tensor.
          mask: optional mask, its application will be fused to the softmax computation if triton is used
          causal: optional performance optimization, if triton is used and the attention is causal
  
      Returns:
          a Tensor of the same dimension and shape as the input with
          values in the range [0, 1] and sum to 1
      """
      return _softmax_dispatch(x, log=False, mask=mask, causal=causal)
  
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 1 for calling a global function with example usage: softmax(x, mask, causal) and returns: _softmax_dispatch