code: def verify_softmax(shape, axis, dtype="float32", log_softmax=False):
      cudnn_op = cudnn.log_softmax if log_softmax else cudnn.softmax
      testing_op = (
          tvm.topi.testing.log_softmax_python
          if log_softmax
          else tvm.topi.testing.softmax_python
      )
  
      A = te.placeholder(shape, dtype=dtype, name="A")
      B = cudnn_op(A, axis)
      s = te.create_schedule([B.op])
  
      dev = tvm.cuda(0)
      a_np = np.random.uniform(size=shape).astype(dtype)
      b_np = testing_op(a_np)
      a = tvm.nd.array(a_np, dev)
      b = tvm.nd.array(b_np, dev)
      f = tvm.build(
          s, [A, B], target="cuda --host=llvm", name="softmax"
      )
      f(a, b)
      tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-3)
  
imports: import sys
  import pytest
  import tvm
  import tvm.testing
  from tvm import te
  from tvm import relay
  from tvm.contrib import cudnn
  from tvm.contrib.nvcc import have_fp16
  from tvm.contrib import graph_executor
  import numpy
  import tvm.topi.testing
  import tvm.testing
  from tvm.relay.op.contrib.cudnn import partition_for_cudnn
label_desc: Use image node_id 5 for calling a global function with example usage: verify_softmax(shape, axis, dtype, log_softmax) without return types