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