@triton.jit
def triton_cross_entropy_backward(input_grad_ptr, input_stride,
    grad_output_ptr, num_classes, BLOCK_SIZE: tl.constexpr):
    row_id = tl.program_id(0).to(tl.int64)
    input_grad_ptr += row_id * input_stride
    grad_output = tl.load(grad_output_ptr)
    for i in range(0, num_classes, BLOCK_SIZE):
        input_offsets = i + tl.arange(0, BLOCK_SIZE)
        input_grad_block = tl.load(input_grad_ptr + input_offsets, mask=
            input_offsets < num_classes)
        tl.store(input_grad_ptr + input_offsets, input_grad_block *
            grad_output, mask=input_offsets < num_classes)
