def test_recursion(): """ Program: let f(n: i32, data: f32) -> f32 = { if (n == 0) { return data; } else { return f(n - 1, log(data)); } } f(2, 10000); """ f = relay.Var("f") f1 = relay.Var("f1") n = relay.Var("n", e.int32) data = relay.Var("data", e.float32) funcbody = relay.If(equal(n, relay.const(0)), data, relay.Call(f1, [subtract(n, relay.const(1)), log(data)])) value = relay.Function([n, data], funcbody, e.float32, []) orig = relay.Let(f, value, relay.Call(f, [relay.const(2), relay.const(10000.0)])) dced = transform.OptimizeOnExpr(orig, transform.DeadCodeElimination()) orig = transform.OptimizeOnExpr(orig, transform.InferType()) assert graph_equal(dced, orig) dced = transform.OptimizeOnExpr(relay.Let(f, value, e.three), transform.DeadCodeElimination()) assert alpha_equal(dced, e.three)
def test_tuple_get_item(): tt = relay.TupleType([e.float32, e.float32]) t = relay.Var('t', tt) a = relay.Var('a') g = relay.TupleGetItem(t, 0) dced = transform.OptimizeOnExpr(g, transform.DeadCodeElimination()) assert alpha_equal(Function(free_vars(dced), dced), Function(free_vars(g), g)) orig = relay.TupleGetItem(relay.Let(a, e.one, t), 0) dced = transform.OptimizeOnExpr(orig, transform.DeadCodeElimination()) assert alpha_equal(Function(free_vars(dced), dced), Function(free_vars(g), g))
def test_round_trip(): x = relay.Var('x') y = relay.Var('y') z = relay.Var('z') body = relay.Let(z, op.add(y, y), op.add(z, z)) body = relay.Let(y, op.add(x, x), body) f = relay.Function([], relay.Let(x, relay.const(1), body)) g = transform.OptimizeOnExpr(f, transform.ToGraphNormalForm()) h = transform.OptimizeOnExpr(g, transform.ToANormalForm()) assert Feature.fLet in detect_feature(f) assert not Feature.fLet in detect_feature(g) check_eval(f, [], 8.0) check_eval(g, [], 8.0) check_eval(h, [], 8.0)
def test_tuple(): t = TypeVar("t") x = Var("x", t) body = TupleGetItem(relay.Tuple([relay.const(4.0), x]), 1) f = Function([x], body, None, [t]) expected = relay.Function([x], x, None, [t]) expected = transform.OptimizeOnExpr(expected, transform.InferType()) assert alpha_equal(dcpe(f), expected)
def test_let(): x = relay.Var("x") y = relay.Var("y") d = relay.const(4.0, 'float32') body = relay.Let(y, x, x + y) body = relay.Let(x, d, body) check_eval(body, 8) opt_body = transform.OptimizeOnExpr(body, transform.ToANormalForm()) check_eval(opt_body, 8)
def test_if(): cond = relay.const(True) x = relay.If(cond, relay.const(2), relay.const(3)) anf = transform.OptimizeOnExpr( x, [transform.ToANormalForm(), transform.InferType()]) a = relay.Var('a', relay.IncompleteType()) b = relay.Var('b', relay.IncompleteType()) c = relay.Var('c', relay.IncompleteType()) d = relay.Var('d', relay.IncompleteType()) true_branch = relay.Let(a, relay.const(2), a) false_branch = relay.Let(b, relay.const(3), b) expected_output = relay.If(c, true_branch, false_branch) expected_output = relay.Let(d, expected_output, d) expected_output = relay.Let(c, cond, expected_output) expected_output = transform.OptimizeOnExpr(expected_output, transform.InferType()) assert alpha_equal(anf, expected_output)
def test_function(): t = relay.TensorType((), 'float32') x = relay.Var("x", t) f = relay.Function([x], x + x) d = relay.const(4.0, 'float32') anf_f = transform.OptimizeOnExpr(f, transform.ToANormalForm()) assert isinstance(anf_f, relay.Function) check_eval(f(d), 8) check_eval(anf_f(d), 8)
def test_empty_ad(): shape = (10, 10) dtype = "float32" t = TensorType(shape, dtype) d = Var("d", t) f = Function([d], d) g = dcpe(f, grad=True) expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])])) expected = transform.OptimizeOnExpr(expected, transform.InferType()) assert alpha_equal(g, expected)
def test_explicit_bound(): x = relay.const(1) y = op.add(x, x) z = op.add(y, y) f = relay.Function([], op.add(z, z)) assert not Feature.fLet in detect_feature(f) anf = transform.OptimizeOnExpr(f, transform.ToANormalForm()) assert Feature.fLet in detect_feature(anf) check_eval(f(), 8.0) check_eval(anf(), 8.0)
def test_ref(): t = relay.TensorType([], "float32") d = relay.Var("d", t) r = relay.Var("r", relay.RefType(t)) x = relay.Var("x") body = relay.RefRead(r) body = Let(x, RefWrite(r, RefRead(r) * RefRead(r)), body) body = Let(r, RefCreate(d), body) square = Function([d], body) expected = transform.OptimizeOnExpr(Function([d], d * d), transform.InferType()) assert alpha_equal(dcpe(square), expected)
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 dcpe(expr, mod=None, grad=False): passes = [ transform.PartialEvaluate(), transform.DeadCodeElimination(inline_once=True) ] if grad: expr = gradient(expr) if mod: assert isinstance(expr, Function) mod[mod.entry_func] = expr seq = transform.Sequential(passes) mod = seq(mod) return mod[mod.entry_func] return transform.OptimizeOnExpr(expr, passes)
def test_order(): z = relay.const(3) y = relay.const(2) x = relay.const(1) val = x + y * z check_eval(val, 7.0) anf = transform.OptimizeOnExpr( val, [transform.ToANormalForm(), transform.InferType()]) a = relay.Var('a', relay.IncompleteType()) b = relay.Var('b', relay.IncompleteType()) c = relay.Var('c', relay.IncompleteType()) d = relay.Var('d', relay.IncompleteType()) e = relay.Var('e', relay.IncompleteType()) expected_output = e expected_output = relay.Let(e, a + d, expected_output) expected_output = relay.Let(d, b * c, expected_output) expected_output = relay.Let(c, z, expected_output) expected_output = relay.Let(b, y, expected_output) expected_output = relay.Let(a, x, expected_output) expected_output = transform.OptimizeOnExpr(expected_output, transform.InferType()) assert alpha_equal(anf, expected_output)
def test_ad(): shape = (10, 10) dtype = "float32" t = TensorType(shape, dtype) d = Var("d", t) f = Function([d], d * d) g = dcpe(f, grad=True) m = d * d x = relay.Var("x") o = op.ones_like(x) x1 = relay.Var("x1") grad = op.zeros_like(d) + op.collapse_sum_like( x1 * d, d) + op.collapse_sum_like(x1 * d, d) body = Tuple([x, Tuple([grad])]) body = relay.Let(x1, o, body) expected = Function([d], relay.Let(x, m, body)) expected = transform.OptimizeOnExpr(expected, transform.InferType()) assert alpha_equal(g, expected)
def test_op_let(): dced = transform.OptimizeOnExpr(add(relay.Let(e.a, e.one, e.three), e.two), transform.DeadCodeElimination()) assert alpha_equal(dced, add(e.three, e.two))
def test_let(): orig = relay.Let(e.x, e.y, e.z) orig = transform.OptimizeOnExpr(orig, transform.DeadCodeElimination()) assert alpha_equal(Function(free_vars(orig), orig), Function([e.z], e.z))
def tipe(expr): return transform.OptimizeOnExpr(expr, [ transform.InferType(), transform.PartialEvaluate(), transform.InferType() ])
def test_used_let(): orig = relay.Let(e.c, e.one, e.c + e.c) orig = transform.OptimizeOnExpr(orig, transform.DeadCodeElimination()) expected = relay.Let(e.c, e.one, e.c + e.c) assert alpha_equal(Function([e.c], orig), Function([e.c], expected))
def test_inline(): orig = relay.Let(e.a, e.b, relay.Let(e.c, e.d, e.c)) orig = transform.OptimizeOnExpr(orig, transform.DeadCodeElimination()) assert alpha_equal(Function(free_vars(orig), orig), Function([e.d], e.d))
def test_chain_unused_let(): orig = relay.Let(e.a, e.b, relay.Let(e.c, e.d, e.e)) orig = transform.OptimizeOnExpr(orig, transform.DeadCodeElimination()) assert alpha_equal(Function(free_vars(orig), orig), Function([e.e], e.e))