def test_ref():
i = relay.Var("i")
iv = relay.Var("iv")
u = relay.Var("u")
uv = relay.Var("uv")
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
check_eval(body, 3)
opt_body = run_opt_pass(body, transform.ToANormalForm())
check_eval(opt_body, 3)
def test_ref():
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
check_eval(body, 3)
opt_body = transform.OptimizeOnExpr(body, transform.ToANormalForm())
check_eval(opt_body, 3)
def test_ref_execution_order():
# we want to have effects execute from left to right
x = relay.Var('x')
y = relay.Var('y')
f = relay.Var('f')
r = relay.Var('r')
expr = relay.Let(f, relay.Function([x, y], x),
# r = 1
relay.Let(r, relay.RefCreate(relay.const(1)),
relay.Tuple([
# should be 1
relay.RefRead(r),
# set r to 2 and read back
seq(relay.RefWrite(r, relay.const(2)),
relay.RefRead(r)),
# set r to 3 and read back
seq(relay.RefWrite(r, relay.const(3)),
relay.RefRead(r)),
# set r to 4 and read as first arg to f
# set r to 5 and read as second arg to f
# f should evaluate to 4
f(
seq(relay.RefWrite(r, relay.const(4)),
relay.RefRead(r)),
seq(relay.RefWrite(r, relay.const(5)),
relay.RefRead(r))),
# read back 5
relay.RefRead(r)
])))
tup_val = run_as_python(expr)
assert_tuple_value(tup_val, 5)
assert_tensor_value(tup_val.fields[0], 1)
assert_tensor_value(tup_val.fields[1], 2)
assert_tensor_value(tup_val.fields[2], 3)
assert_tensor_value(tup_val.fields[3], 4)
assert_tensor_value(tup_val.fields[4], 5)
def test_ref():
mod = relay.Module()
three_with_ref = relay.GlobalVar('three_with_ref')
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
mod[three_with_ref] = relay.Function([], body)
check_eval(three_with_ref, [], 3, mod=mod)
def test_ref():
mod = tvm.IRModule()
three_with_ref = relay.GlobalVar("three_with_ref")
i = relay.Var("i")
iv = relay.Var("iv")
u = relay.Var("u")
uv = relay.Var("uv")
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
mod[three_with_ref] = relay.Function([], body)
check_eval(three_with_ref, [], 3, mod=mod)
def test_ref():
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
check_eval(body, 3)
opt_body = run_opt_pass(body, transform.ToBasicBlockNormalForm())
check_eval(opt_body, 3)
check_basic_block_normal_form(opt_body)
def test_ref_write():
# check that the result of a ref write is an empty tuple
v = relay.Var('v')
initial_write = relay.Let(v, relay.RefCreate(relay.Tuple([relay.const(1)])),
relay.RefWrite(v, relay.Tuple([relay.const(2)])))
write_val = run_as_python(initial_write)
assert_tuple_value(write_val, 0)
# now ensure that the value, once written, can be read back
# (we read the value before and after mutation)
w = relay.Var('w')
read_after_write = relay.Let(
v, relay.RefCreate(relay.Tuple([relay.const(1)])),
relay.Let(
w, relay.RefCreate(relay.RefRead(v)),
seq(relay.RefWrite(v, relay.Tuple([relay.const(2)])),
relay.Tuple([relay.RefRead(w), relay.RefRead(v)]))))
read_val = run_as_python(read_after_write)
assert_tuple_value(read_val, 2)
assert_tuple_value(read_val.fields[0], 1)
assert_tuple_value(read_val.fields[1], 1)
assert_tensor_value(read_val.fields[0].fields[0], 1)
assert_tensor_value(read_val.fields[1].fields[0], 2)
def test_ref():
mod = relay.Module()
three_with_ref = relay.GlobalVar('three_with_ref')
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2, dtype='int32')), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1, dtype='int32')), body)
mod[three_with_ref] = relay.Function([], body)
cfunc = compile(three_with_ref, mod)
output = cfunc()
np.testing.assert_allclose(output.asnumpy(), np.array(3, dtype='int32'))
def test_ref():
shape = (10, 10)
dtype = 'float32'
t = relay.TensorType(shape, dtype)
x = relay.var("x", t)
r = relay.Var("r")
u = relay.Var("u")
body = relay.RefRead(r)
body = relay.Let(u, relay.RefWrite(r, relay.RefRead(r) + relay.RefRead(r)), body)
body = relay.Let(r, relay.RefCreate(x), body)
func = relay.Function([x], body)
back_func = relay.ir_pass.infer_type(gradient(func))
assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])]))
x_nd = rand(dtype, *shape)
ex = create_executor()
forward, (grad_x,) = ex.evaluate(back_func)(x_nd)
tvm.testing.assert_allclose(forward.asnumpy(), 2 * x_nd.asnumpy())
tvm.testing.assert_allclose(grad_x.asnumpy(), 2 * np.ones_like(grad_x.asnumpy()))
def test_if_ref():
shape = ()
dtype = "bool"
t = relay.TensorType(shape, dtype)
d = relay.Var("d", t)
r = relay.Var("r")
update = relay.Function([], relay.RefWrite(r, relay.RefRead(r) + relay.RefRead(r)))
u = relay.Var("u")
body = relay.If(d, u(), u())
eff = relay.Var("eff")
body = relay.Let(eff, body, relay.RefRead(r))
f = relay.Function([d], relay.Let(r, relay.RefCreate(relay.const(1)), relay.Let(u, update, body)))
f = infer_type(f)
pe_f = infer_type(partial_evaluate(f))
ex = create_executor()
f_res = ex.evaluate(f)(relay.const(True))
pe_f_res = ex.evaluate(pe_f)(relay.const(True))
np.testing.assert_allclose(f_res.asnumpy(), 2 * np.ones_like(f_res.asnumpy()))
np.testing.assert_allclose(pe_f_res.asnumpy(), 2 * np.ones_like(pe_f_res.asnumpy()))
def test_arbitrary_let_nesting():
# something that is tricky to do in Python but comes naturally in Relay
mod = relay.Module()
p = Prelude(mod)
x = relay.Var('x')
r = relay.Var('r')
y = relay.Var('y')
z = relay.Var('z')
expr = relay.Tuple([
relay.Let(x, relay.Tuple([relay.const(1),
relay.const(2)]), relay.TupleGetItem(x, 1)),
relay.Let(r, relay.RefCreate(relay.const(1)),
seq(relay.RefWrite(r, relay.const(3)), relay.RefRead(r))),
relay.Let(y, p.id(relay.Let(z, relay.const(4), z)), y)
])
tup_val = run_as_python(expr, mod)
assert_adt_len(tup_val, 3)
assert_tensor_value(tup_val[0], 2)
assert_tensor_value(tup_val[1], 3)
assert_tensor_value(tup_val[2], 4)
def test_match_effect_exactly_once():
mod = relay.Module()
p = Prelude(mod)
# the list should be of length 1!
# Unless we mistakenly execute the data clause more than once
r = relay.Var('r')
data = seq(relay.RefWrite(r, p.cons(relay.Tuple([]), relay.RefRead(r))), relay.RefRead(r))
match = relay.Let(
r, relay.RefCreate(p.nil()),
relay.Match(data, [
relay.Clause(relay.PatternConstructor(p.nil, []), relay.const(0)),
relay.Clause(
relay.PatternConstructor(
p.cons,
[relay.PatternWildcard(), relay.PatternConstructor(p.nil, [])]),
relay.const(1)),
relay.Clause(relay.PatternWildcard(), relay.const(2))
]))
match_val = run_as_python(match, mod)
assert_tensor_value(match_val, 1)
def test_function_invalidate():
shape = ()
dtype = "bool"
t = relay.TensorType(shape, dtype)
d = relay.Var("d", t)
r = relay.Var("r")
fetch = relay.Function([], relay.RefRead(r))
fet = relay.Var("fetch")
fet_obscured = relay.Var("fetch_obscured")
u = relay.Var("u")
body = relay.If(d, fet_obscured(), fet_obscured())
body = relay.Let(u, relay.RefWrite(r, relay.const(1)), body)
body = relay.Let(fet_obscured, relay.If(d, fet, fet), body)
body = relay.Let(fet, fetch, body)
body = relay.Let(r, relay.RefCreate(relay.const(0)), body)
f = relay.Function([d], body)
f = infer_type(f)
pe_f = infer_type(partial_evaluate(f))
ex = create_executor()
f_res = ex.evaluate(f)(relay.const(True))
pe_f_res = ex.evaluate(pe_f)(relay.const(True))
np.testing.assert_allclose(f_res.asnumpy(), np.ones_like(f_res.asnumpy()))
np.testing.assert_allclose(pe_f_res.asnumpy(), np.ones_like(pe_f_res.asnumpy()))
def relay_universe_setitem(c, u, h, v):
return relay.RefWrite(c.ref(h), c.ref(v))
