def METHOD_NAME(shape, axis, dtype="float32", log_softmax=False):
    miopen_op = miopen.log_softmax if log_softmax else miopen.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 = miopen_op(A, axis)
    s = te.create_schedule([B.op])
    dev = tvm.rocm(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="rocm --host=llvm", name="softmax")
    f(a, b)
    tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-3)