def make_matrix_softmax_cross_entropy(shape,
tgt,
tgt_host,
func_name,
dtype="float32"):
"""Hint: output shape should be (1,)"""
y = tvm.placeholder(shape, dtype=dtype, name="y")
y_ = tvm.placeholder(shape, dtype=dtype, name="y_")
t = -topi.sum(y_ * topi.log(topi.nn.softmax(y)), axis=1)
c = topi.sum(t, keepdims=True) / shape[0]
s = tvm.create_schedule(c.op)
f = tvm.build(s, [y, y_, c], tgt, target_host=tgt_host, name=func_name)
return f
def log_compute(attrs, inputs, output_type, target):
assert len(inputs) == 1
return [topi.log(inputs[0])]
def compute_cross_entropy(attrs, inputs, out_dtype):
x, y = inputs
return [-topi.sum(topi.log(x) * y) / x.shape[0]]
