def test_where_simple():
data = relay.var("data", shape=[1, 20])
weight = relay.var("weight", shape=[20, 20])
a = relay.nn.dense(data, weight, units=20)
b = relay.where(data, a, a)
mod = tvm.IRModule.from_expr(b)
mod_params = {
"data": np.random.uniform(-1, 1, size=[1, 20]).astype("float32"),
"weight": np.random.uniform(-1, 1, size=[20, 20]).astype("float32"),
}
output_mod = verify_mixed_precision_output_close(mod,
mod_params,
atol=0.01,
rtol=0.01)
# Create expected module
data = relay.cast(relay.var("data", shape=[1, 20]), "float16")
weight = relay.cast(relay.var("weight", shape=[20, 20]), "float16")
a = relay.nn.dense(data, weight, units=20, out_dtype="float16")
b = relay.where(data, a, a)
expected_mod = tvm.IRModule.from_expr(b)
expected_mod = InferType()(expected_mod)
assert tvm.ir.structural_equal(expected_mod, output_mod)
def verify_any_where(cond_shape,
x_shape,
y_shape,
cond_np_shape,
x_np_shape,
y_np_shape,
y_np_shape_invalid=None):
dtype = "float32"
cond = relay.var("cond", shape=cond_shape, dtype="bool")
x = relay.var("x", shape=x_shape, dtype=dtype)
y = relay.var("y", shape=y_shape, dtype=dtype)
z = relay.where(cond, x, y)
mod = tvm.IRModule()
mod["main"] = relay.Function([cond, x, y], z)
cond_np = np.random.randn(*cond_np_shape) > 0
x_np = np.random.randn(*x_np_shape).astype(dtype)
y_np = np.random.randn(*y_np_shape).astype(dtype)
expected = np.where(cond_np, x_np, y_np)
check_result([cond_np, x_np, y_np], mod, expected)
# verify invalid broadcasting check
if y_np_shape_invalid:
y_np_bad = np.random.randn(*y_np_shape_invalid).astype(dtype)
try:
check_result([cond_np, x_np, y_np_bad], mod, expected)
except tvm.error.TVMError as e:
error_msg = str(e).split("\n")[-1]
assert "Invalid broadcast shapes" in error_msg
def verify(x_np, y_np, cond_np):
ref_res = np.where(cond_np, x_np, y_np)
args = []
args_np = []
vs = []
cond = relay.var("cond", relay.TensorType(cond_np.shape, "bool"))
args.append(cond)
args_np.append(cond_np)
for v_name, v_np in [("x", x_np), ("y", y_np)]:
if len(v_np.shape) == 0:
v = relay.const(v_np.item())
else:
v = relay.var(v_name, relay.TensorType(v_np.shape, dtype))
args.append(v)
args_np.append(v_np)
vs.append(v)
z = relay.where(cond, vs[0], vs[1])
func = relay.Function(args, z)
run(func, args_np, ref_res)
def test_where():
cond = relay.var("cond", relay.TensorType((3, 4), "float32"))
x = relay.var("x", relay.TensorType((3, 4), "float32"))
y = relay.var("y", relay.TensorType((3, 4), "float32"))
z = relay.where(cond, x, y)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.TensorType((3, 4), "float32")
def test_where():
cond = relay.var("cond", relay.TensorType((3, 4), "float32"))
x = relay.var("x", relay.TensorType((3, 4), "float32"))
y = relay.var("y", relay.TensorType((3, 4), "float32"))
z = relay.where(cond, x, y)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.TensorType((3, 4), "float32")
def test_where():
""" from test_op_level4.py, added debug printing """
cond = relay.var("cond", relay.TensorType((3, 4), "float32"))
x = relay.var("x", relay.TensorType((3, 4), "float32"))
y = relay.var("y", relay.TensorType((3, 4), "float32"))
z = relay.where(cond, x, y)
print(z.astext())
def safe_exp(w):
slope = relay.const(np.exp(1, dtype=np.float32))
lin_bool = w > slope
lin_region = relay.cast(lin_bool, "float32")
lin_out = slope * w
exp_out = relay.exp(relay.where(lin_bool, relay.zeros_like(w), w))
out = lin_region * lin_out + (relay.const(1.) - lin_region) * exp_out
return out
def test_where():
ib = relay.ir_builder.IRBuilder()
cond = ib.param("cond", relay.TensorType((3, 4), "float32"))
x = ib.param("x", relay.TensorType((3, 4), "float32"))
y = ib.param("y", relay.TensorType((3, 4), "float32"))
with ib.function(cond, x, y) as func:
ib.ret(relay.where(cond, x, y))
ib.ret(func)
func = relay.ir_pass.infer_type(ib.env, func.to_func())
ftype = func.checked_type
assert ftype.ret_type == relay.TensorType((3, 4), "float32")
def test_where():
def run(func, inputs, ref_res):
for target, ctx in tvm.testing.enabled_targets():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(*inputs)
tvm.testing.assert_allclose(op_res.asnumpy(),
ref_res,
rtol=1e-5)
shape = (3, 4)
dtype = "float32"
cond = relay.var("cond", relay.TensorType(shape, dtype))
x = relay.var("x", relay.TensorType(shape, dtype))
y = relay.var("y", relay.TensorType(shape, dtype))
z = relay.where(cond, x, y)
zz = run_infer_type(z)
assert zz.checked_type == relay.TensorType(shape, dtype)
func = relay.Function([cond, x, y], z)
condition = np.random.uniform(low=-1, high=1, size=shape).astype(dtype)
x = np.random.uniform(size=shape).astype(dtype)
y = np.random.uniform(size=shape).astype(dtype)
ref_res = np.where(condition, x, y)
run(func, [condition, x, y], ref_res)
x = relay.const(1)
y = relay.const(-1)
shape = (3, )
dtype = "float32"
cond = relay.var("cond", relay.TensorType(shape, "bool"))
z = relay.where(cond, x, y)
func = relay.Function([cond], z)
condition = np.array([1, 0, 1], dtype=np.bool)
ref_res = np.where(condition, 1, -1)
run(func, [condition], ref_res)
def test_where_constant(self):
x = relay.var("x", relay.TensorType((-1, 2, 2, 4), "float32"))
y = relay.var("y", relay.TensorType((-1, 2, 2, 4), "float32"))
c = relay.expr.const(np.ones((4, ), np.float32))
w = relay.where(c, x, y)
net = relay.Function([x, y], w)
mod = tvm.IRModule.from_expr(net)
mod = relay.transform.InferType()(mod)
xgraph = xf_relay.from_relay(mod, {})
layers = xgraph.get_layers()
assert layers[0].type[0] == "Constant"
assert layers[1].type[0] == "Input"
assert layers[2].type[0] == "Input"
assert layers[3].type[0] == "AnyOp"
assert layers[3].shapes == [-1, 2, 2, 4]
def test_where_grad():
cond_type = relay.TensorType((2, 3, 4), "int32")
lhs_type = relay.TensorType((1, 3, 4), "float32")
rhs_type = relay.TensorType((2, 1, 4), "float32")
inputs = [
np.random.randint(2,
size=cond_type.concrete_shape,
dtype=cond_type.dtype),
_np_randn_from_type(lhs_type, scale=1e-5),
_np_randn_from_type(rhs_type, scale=1e-5),
]
cond = relay.var("cond", type_annotation=cond_type)
lhs = relay.var("lhs", type_annotation=lhs_type)
rhs = relay.var("rhs", type_annotation=rhs_type)
fwd_func = relay.Function([cond, lhs, rhs], relay.where(cond, lhs, rhs))
check_grad(fwd_func, inputs=inputs, test_inputs=inputs[1:])
def test_where():
shape = (3, 4)
dtype = "float32"
cond = relay.var("cond", relay.TensorType(shape, dtype))
x = relay.var("x", relay.TensorType(shape, dtype))
y = relay.var("y", relay.TensorType(shape, dtype))
z = relay.where(cond, x, y)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.TensorType(shape, dtype)
func = relay.Function([cond, x, y], z)
condition = np.random.uniform(low=-1, high=1, size=shape).astype(dtype)
x = np.random.uniform(size=shape).astype(dtype)
y = np.random.uniform(size=shape).astype(dtype)
ref_res = np.where(condition, x, y)
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)(condition, x, y)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def test_where():
shape = (3, 4)
dtype = "float32"
cond = relay.var("cond", relay.TensorType(shape, dtype))
x = relay.var("x", relay.TensorType(shape, dtype))
y = relay.var("y", relay.TensorType(shape, dtype))
z = relay.where(cond, x, y)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.TensorType(shape, dtype)
func = relay.Function([cond, x, y], z)
condition = np.random.uniform(low=-1, high=1, size=shape).astype(dtype)
x = np.random.uniform(size=shape).astype(dtype)
y = np.random.uniform(size=shape).astype(dtype)
ref_res = np.where(condition, x, y)
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)(condition, x, y)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
