github-featured/apache--tvm
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Relay Op sprint (part 2) - Level 1 - log_softmax (#2128)

Browse Source
This commit is contained in:
Animesh Jain
2018-11-19 14:23:37 -05:00
committed by Tianqi Chen
parent 81da33f81b
commit b71edd76bd
3 changed files with 32 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relay/op/nn/_nn.py
+9 -1
View File
@@ -9,7 +9,6 @@ from ..op import OpPattern, schedule_injective
reg.register_schedule("nn.relu", schedule_injective)
reg.register_pattern("nn.relu", OpPattern.ELEMWISE)
@reg.register_schedule("nn.softmax")
def schedule_softmax(_, outputs, target):
"""Schedule definition of softmax"""
@@ -19,6 +18,15 @@ def schedule_softmax(_, outputs, target):
reg.register_pattern("nn.softmax", OpPattern.OPAQUE)
@reg.register_schedule("nn.log_softmax")
def schedule_log_softmax(_, outputs, target):
"""Schedule definition of log_softmax"""
with target:
return topi.generic.schedule_softmax(outputs)
reg.register_pattern("nn.log_softmax", OpPattern.OPAQUE)
# dense
@reg.register_compute("nn.dense")
def compute_dense(attrs, inputs, out_type, target):