code: def test_op(M, N, dtype, mode):
      capability = torch.cuda.get_device_capability()
      if capability[0] < 8 and dtype == "bfloat16":
          pytest.skip("Only test bfloat16 on devices with sm >= 80")
      dtype = {
          "bfloat16": torch.bfloat16,
          "float16": torch.float16,
          "float32": torch.float32,
      }[dtype]
      # create inputs
      x = torch.randn(
          M, N, dtype=dtype, device="cuda", requires_grad=True
      )
      idx = 4 + torch.ones(M, dtype=torch.int64, device="cuda")
      # forward pass
      tt_y = triton.ops.cross_entropy(x, idx)
      th_y = torch.nn.CrossEntropyLoss(reduction="none")(x, idx)
      if mode == "forward":
          torch.testing.assert_close(th_y, tt_y)
      # backward pass
      elif mode == "backward":
          dy = torch.randn_like(tt_y)
          # triton backward
          tt_y.backward(dy)
          tt_dx = x.grad.clone()
          # torch backward
          x.grad = None
          th_y.backward(dy)
          th_dx = x.grad.clone()
          if dtype == torch.float16:
              torch.testing.assert_close(
                  th_dx, tt_dx, rtol=0.001, atol=0.001
              )
          else:
              torch.testing.assert_close(th_dx, tt_dx)
  
imports: import pytest
  import torch
  import triton
  import triton.ops
label_desc: Use image node_id 1 for calling a global function with example usage: test_op(M, N, dtype, mode) without return types