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)
