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(): """ 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)