code: def backward(ctx, grad_losses):
      logits, lse, labels = ctx.saved_tensors
      dlogits = (
          logits if ctx.inplace_backward else torch.empty_like(logits)
      )
      n_rows, n_cols = logits.shape
      BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024)
      num_warps = (
          4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16)
      )
      grid = lambda META: (
          n_rows,
          triton.cdiv(n_cols, META["BLOCK_SIZE"]),
      )  # noqa
      # Need this, otherwise Triton tries to launch from cuda:0 and we get
      # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
      with torch.cuda.device(logits.device.index):
          cross_entropy_bwd_kernel[grid](
              dlogits,  # data ptrs
              grad_losses,
              logits,
              lse,
              labels,
              ctx.smoothing,
              ctx.logit_scale,
              ctx.lse_square_scale,
              ctx.ignored_index,
              ctx.total_classes,
              ctx.class_start_idx,
              n_cols,  # shapes
              logits.stride(0),  # strides
              dlogits.stride(0),
              grad_losses.stride(0),
              BLOCK_SIZE=BLOCK_SIZE,  # constants
              num_warps=num_warps,
          )
      return dlogits, None, None, None, None, None, None, None
  
file_path: flash-attention/flash_attn/ops/triton/cross_entropy.py
imports: from typing import Tuple, Optional, Union
  import torch
  from einops import rearrange
  import triton
  import triton.language
label_desc: Use image node_id 2 for calling the CrossEntropyLoss obj's underlying member method code with example usage: obj.backward(ctx, grad_losses) and returns: dlogits,  None,  None,  None,  None,  None,  None,  None