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]]