code: def benchmark(M, N, provider):
      x = torch.randn(M, N, device="cuda", dtype=torch.float32)
      quantiles = [0.5, 0.2, 0.8]
      if provider == "torch-native":
          ms, min_ms, max_ms = triton.testing.do_bench(
              lambda: torch.softmax(x, axis=-1), quantiles=quantiles
          )
      if provider == "triton":
          ms, min_ms, max_ms = triton.testing.do_bench(
              lambda: softmax(x), quantiles=quantiles
          )
      if provider == "torch-jit":
          ms, min_ms, max_ms = triton.testing.do_bench(
              lambda: naive_softmax(x), quantiles=quantiles
          )
      gbps = (
          lambda ms: 2
          * x.nelement()
          * x.element_size()
          * 1e-9
          / (ms * 1e-3)
      )
      return gbps(ms), gbps(max_ms), gbps(min_ms)
  
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 4 for calling a global function with example usage: benchmark(M, N, provider) and returns: gbps,  gbps,  gbps