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))