def verify_arange(start, stop, step): if start is None and step is None: A = topi.arange(stop) a_np = np.arange(stop) elif start is None: A = topi.arange(stop, step=step) a_np = np.arange(stop, step=step) elif step is None: A = topi.arange(start, stop) a_np = np.arange(start, stop) else: A = topi.arange(start, stop, step) a_np = np.arange(start, stop, step) def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): s = topi.generic.schedule_injective(A) f = tvm.build(s, [A], device, name="arange") a_nd = tvm.nd.empty(a_np.shape, dtype='float32', ctx=ctx) f(a_nd) tvm.testing.assert_allclose(a_nd.asnumpy(), a_np) for device in get_all_backend(): check_device(device)