code: def _forward(LOGITS, PROBS, IDX, LOSS, N, BLOCK: tl.constexpr):
      row = tl.program_id(0)
      cols = tl.arange(0, BLOCK)
      idx = tl.load(IDX + row)
      # pointers to logit and probs
      LOGITS = LOGITS + row * N + cols
      WRIT_PROBS = PROBS + row * N + cols
      READ_PROBS = PROBS + row * N + idx
      # write-back negative log-probs
      logits = tl.load(LOGITS, mask=cols < N, other=-float("inf"))
      logits = logits.to(tl.float32)
      logits = logits - tl.max(logits, 0)
      probs = tl.log(tl.sum(tl.exp(logits), 0)) - logits
      tl.store(WRIT_PROBS, probs, mask=cols < N)
      # There is a bug in the compiler, which fails to insert a barrier here.
      # We add it explicitly for now. Will be fixed soon.
      tl.debug_barrier()
      # write-back loss
      probs = tl.load(READ_PROBS)
      tl.store(LOSS + row, probs)
  
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 2 for calling a global function with example usage: _forward(LOGITS, PROBS, IDX, LOSS, N, BLOCK) without return types