code: def _softmax_backward(
      GradIn,
      GradOut,
      Out,
      stride_bm,
      stride_bn,
      stride_gm,
      stride_gn,
      stride_om,
      stride_on,
      K,
      # meta-params
      depth: tl.constexpr,
      causal: tl.constexpr,
      log: tl.constexpr,
  ):
      # fmt: on
  
      """
      Compute the softmax gradients.
      ..Note: Not autotuning for now because this would lead to broken accumulated gradients
      """
  
      m = tl.program_id(0)
      n = tl.program_id(1)
  
      # col indices
      k = tl.arange(0, depth)
  
      # the memory address of all the elements that we want to load can be computed as follows
      grad_out_ptrs = GradOut + m * stride_gm + n * stride_gn + k
      out_ptrs = Out + m * stride_om + n * stride_on + k
  
      # load input data; pad out-of-bounds elements with 0
      io_mask = k < K
  
      # Causal - 1: skip on the loads directly
      if causal:
          io_mask = io_mask & (k <= n)
  
      g = tl.load(grad_out_ptrs, mask=io_mask, other=float(0)).to(
          tl.float32
      )
      o = tl.load(out_ptrs, mask=io_mask, other=float(0)).to(tl.float32)
  
      # Causal - 2: enforce correctness over a couple of misloaded values
      if causal:
          zero = float(0)
          zero = zero.to(g.dtype)  # type: ignore
          g = tl.where(k > n, zero, g)
          o = tl.where(k > n, zero, o)
  
      if log:
          s = tl.sum(g, 0)
          grad_in = g - tl.exp(o) * s
      else:
          # Step 1: Compute the intermediate sum used for the gradient
          s = tl.sum(g * o, 0)
  
          # Step 2: Compute the gradients
          grad_in = o * (g - s)
  
      # write back to the input gradients
      # technically we could write only the lower triangular matrix in the causal case
      # but this is deemed to error prone
      grad_in_ptrs = GradIn + m * stride_bm + n * stride_bn + k
      tl.store(grad_in_ptrs, grad_in, mask=k < K)
  
imports: import triton
  import triton.language
label_desc: Use image node_id 3 for calling a global function with example usage: _softmax_backward(GradIn, GradOut, Out, stride_bm, stride_bn, stride_gm, stride_gn, stride_om, stride_on, K, depth, causal, log) without return types