This is a issue with compute registration. Although you declare softmax as generic func, the compute for softmax in Relay is defined in C++ side. So the python generic func is never invoked. See my patch below
diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py
index 54f13c68..9801d4a9 100644
--- a/python/tvm/relay/op/nn/_nn.py
+++ b/python/tvm/relay/op/nn/_nn.py
@@ -31,6 +31,11 @@ reg.register_schedule("nn.relu", schedule_injective)
reg.register_pattern("nn.relu", OpPattern.ELEMWISE)
# softmax
+@reg.register_compute("nn.softmax")
+def compute_softmax(attrs, inputs, out_type, target):
+ """Compute definition of softmax"""
+ return [topi.nn.softmax(inputs[0])]
+
@reg.register_schedule("nn.softmax")
def schedule_softmax(_, outputs, target):
"""Schedule definition of softmax"""
diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc
index d3a71787..5e4f5cfa 100644
--- a/src/relay/op/nn/nn.cc
+++ b/src/relay/op/nn/nn.cc
@@ -348,7 +348,8 @@ RELAY_REGISTER_OP("nn.softmax")
.set_num_inputs(1)
.add_argument("data", "Tensor", "The input tensor.")
.set_support_level(1)
-.add_type_rel("Identity", IdentityRel)
+.add_type_rel("Identity", IdentityRel);
+/*
.set_attr<FTVMCompute>("FTVMCompute", [](const Attrs& attrs,
const Array<Tensor>& inputs,
const Type& out_type,
@@ -357,7 +358,7 @@ RELAY_REGISTER_OP("nn.softmax")
CHECK(param != nullptr);
return Array<Tensor>{ topi::nn::softmax(inputs[0], param->axis) };
});
-
+*/
// relay.nn.log_softmax
TVM_REGISTER_API("relay.op.nn._make.log_softmax")