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 test_cumsum(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.cumsum(x))
x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype)
res_np = np.cumsum(x_np)
check_mod(target, dev, mod, x_np, res_np)
def verify_cumsum(xshape, axis=0):
x = relay.var("x", relay.TensorType(xshape, "float32"))
z = relay.cumsum(x, axis=axis, exclusive=False, reverse=False)
func = relay.Function([x], z)
x_data = np.random.uniform(size=xshape).astype("float32")
if axis < 0:
np_axis = len(xshape) + axis
else:
np_axis = axis
ref_res = np.cumsum(x_data, np_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,
rtol=1e-5)