def check_cumsum(np_ref, data, axis=None, dtype=None):
implementations = {
"generic": (lambda x: topi.cumsum(x, axis, dtype),
topi.generic.schedule_extern),
"cuda": (lambda x: topi.cuda.cumsum(x, axis, dtype),
topi.cuda.schedule_scan),
"nvptx": (lambda x: topi.cuda.cumsum(x, axis, dtype),
topi.cuda.schedule_scan),
}
fcompute, fschedule = tvm.topi.testing.dispatch(
target, implementations)
tvm.topi.testing.compare_numpy_tvm([data], np_ref, target, ctx,
fcompute, fschedule)
def compute_cumsum(attrs, inputs, output_type):
"""Compute definition of cumsum"""
return [topi.cumsum(inputs[0], attrs.axis, attrs.dtype, attrs.exclusive)]
