def verify_argsort(shape, axis, is_ascend, dtype, is_dyn=False): if is_dyn: x = relay.var( "x", relay.TensorType([relay.Any()] * len(shape), "float32")) else: x = relay.var("x", relay.TensorType(shape, "float32")) z = relay.argsort(x, axis=axis, is_ascend=is_ascend, dtype=dtype) func = relay.Function([x], z) x_data = np.random.uniform(size=shape).astype("float32") if is_ascend: ref_res = np.argsort(x_data, axis=axis, kind="stable") else: ref_res = np.argsort(-x_data, axis=axis, kind="stable") if is_dyn: backends = ["vm", "debug"] else: backends = ["graph", "debug"] for target, ctx in tvm.testing.enabled_targets(): for kind in backends: mod = tvm.ir.IRModule.from_expr(func) intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target) op_res = intrp.evaluate()(x_data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res.astype(dtype), rtol=1e-5)
def test_vulkan_pushconstants(): # Three 32 bit pushconstants: any_dim, stride, stride dtype = "float32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.sqrt(x)) x_np = np.random.uniform(size=(10, )).astype(dtype) res_np = np.sqrt(x_np) check_mod(mod, x_np, res_np) # One 64 bit and one 32 bit constants dtype = "int32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.argsort(x)) x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype) res_np = np.argsort(x_np) check_mod(mod, x_np, res_np) # One 64 bit and one 32 bit constants dtype = "int32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.cumsum(x)) x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype) res_np = np.cumsum(x_np) check_mod(mod, x_np, res_np)
def verify_argsort(shape, axis, is_ascend, dtype, is_dyn=False, in_dtype="float32"): if is_dyn: x = relay.var( "x", relay.TensorType([relay.Any()] * len(shape), in_dtype)) else: x = relay.var("x", relay.TensorType(shape, in_dtype)) z = relay.argsort(x, axis=axis, is_ascend=is_ascend, dtype=dtype) func = relay.Function([x], z) x_data = np.random.uniform(size=shape).astype(in_dtype) if is_ascend: ref_res = np.argsort(x_data, axis=axis, kind="stable") else: ref_res = np.argsort(-x_data, axis=axis, kind="stable") if is_dyn: backend = "vm" else: backend = "graph" for target, dev in tvm.testing.enabled_targets(): mod = tvm.ir.IRModule.from_expr(func) op_res = relay.create_executor(backend, mod=mod, device=dev, target=target).evaluate()(x_data) tvm.testing.assert_allclose(op_res.numpy(), ref_res.astype(dtype), rtol=1e-5)
def test_pushconstants(): if not tvm.testing.device_enabled("vulkan"): return def check_mod(mod, x_np, res_np): target = "vulkan" ctx = tvm.context(target, 0) ex = relay.create_executor("vm", mod=mod, ctx=ctx, target=target) res = ex.evaluate()(x_np).asnumpy() tvm.testing.assert_allclose(res, res_np, atol=1e-5) # Three 32 bit pushconstants: any_dim, stride, stride dtype = "float32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.sqrt(x)) x_np = np.random.uniform(size=(10, )).astype(dtype) res_np = np.sqrt(x_np) check_mod(mod, x_np, res_np) # One 64 bit and one 32 bit constants dtype = "int32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.argsort(x)) x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype) res_np = np.argsort(x_np) check_mod(mod, x_np, res_np)
def test_argsort(target, dev): # One 64 bit and one 32 bit constants dtype = "int32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.argsort(x)) x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype) res_np = np.argsort(x_np) check_mod(target, dev, mod, x_np, res_np)
def verify_argsort(shape, axis, is_ascend, dtype): x = relay.var("x", relay.TensorType(shape, "float32")) z = relay.argsort(x, axis=axis, is_ascend=is_ascend, dtype=dtype) func = relay.Function([x], z) x_data = np.random.uniform(size=shape).astype("float32") if is_ascend: ref_res = np.argsort(x_data, axis=axis) else: ref_res = np.argsort(-x_data, axis=axis) for target, ctx in ctx_list(): for kind in ["graph", "debug"]: intrp = relay.create_executor(kind, ctx=ctx, target=target) op_res = intrp.evaluate(func)(x_data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res.astype(dtype), rtol=1e-5)
def verify_argsort(shape, axis, is_ascend): x = relay.var("x", relay.TensorType(shape, "float32")) z = relay.argsort(x, axis=axis, is_ascend=is_ascend) zz = relay.ir_pass.infer_type(z) func = relay.Function([x], z) x_data = np.random.uniform(size=shape).astype("float32") if is_ascend: ref_res = np.argsort(x_data, axis=axis) else: ref_res = np.argsort(-x_data, axis=axis) for target, ctx in ctx_list(): for kind in ["graph", "debug"]: intrp = relay.create_executor(kind, ctx=ctx, target=target) op_res = intrp.evaluate(func)(x_data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res.astype("float"), rtol=1e-5)