code: def _backward(PROBS, IDX, DPROBS, N, BLOCK: tl.constexpr):
      row = tl.program_id(0)
      cols = tl.arange(0, BLOCK)
      idx = tl.load(IDX + row)
      # pointers to probs
      PROBS = PROBS + row * N + cols
      # We know d(-log(p[i])/dlogit[k] = -id_mat[i,k] + p[k]
      # and we have -log(p[k]) stored in PROBS, so this is easy
      probs = -tl.load(PROBS, mask=cols < N, other=float("inf"))
      probs = tl.exp(probs.to(tl.float32))
      delta = cols == idx
      # write result in-place in PROBS
      dout = tl.load(DPROBS + row)
      din = (probs - delta) * dout
      tl.store(PROBS, din.to(PROBS.dtype.element_ty), mask=cols < N)
  
file_path: triton/python/triton/ops/cross_entropy.py
imports: import torch
  from ..None import heuristics, jit
  from ..None import language
  from ..None import next_power_of_2
label_desc: Use image node_id 3 for calling a global function with example usage: _backward(PROBS, IDX, DPROBS, N, BLOCK) without return types