code: def log_softmax(
      x: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
      causal: bool = False,
  ) -> torch.Tensor:
      r"""Applies the :math:`\log(\text{Softmax}(x))` function to an 3-dimensional
      input Tensor. The LogSoftmax formulation can be simplified as:
  
      .. math::
          \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)
  
      Args:
          x: input tensor.
  
      Returns:
          a Tensor of the same dimension and shape as the input with
          values in the range [-inf, 0)
      """
      return _softmax_dispatch(x, log=True, 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 2 for calling a global function with example usage: log_softmax(x, mask, causal) and returns: _softmax_dispatch