def __init__(self, *, name="f", **kwargs):
name = f"{name}_{Network.cnt}"
Network.cnt += 1
if len(Network.stack) is not 0:
mod = Network.stack[-1].mod
p = Network.stack[-1].p
else:
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
self.mod = mod
self.p = p
self.inputs = []
self.weights = OrderedSet()
self.sub_network = OrderedSet()
self.f = relay.GlobalVar(name)
self.recurse = relay.Var("recurse")
self.use_recurse = False
self.ret_type = None
body = self.build(**kwargs)
assert isinstance(body, relay.Expr)
if self.use_recurse:
inputs = [copy_var(v) for v in self.inputs]
body = relay.Let(
self.recurse,
relay.Function(inputs, self.call_from_outside(*inputs)), body)
self.mod[self.f] = relay.Function(self.inputs + self.all_weights(),
body, self.ret_type)
def test_nat_3():
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
cfunc = compile(Function([], p.s(p.s(p.s(p.z())))), mod)
output = cfunc()
assert nat_to_int(output) == 3
def test_double():
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
orig = p.double(make_nat_expr(p, 3))
res = dcpe(orig, mod=mod)
assert alpha_equal(res, make_nat_expr(p, 6))
def test_42():
mod = Module()
func = Function([], relay.const(42))
cfunc = compile(func, mod)
output = cfunc()
np.testing.assert_allclose(output.asnumpy(), np.array(42.0,
dtype='float32'))
def test_add_convert():
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
cfunc = compile(p.add, mod)
output = cfunc(int_to_nat(p, 12), int_to_nat(p, 34))
assert nat_to_int(output) == 46
def test_identity():
mod = Module()
x = var('x', shape=())
func = Function([x], x)
cfunc = compile(func, mod)
a = tvm.nd.array(np.array(1.0, dtype='float32'))
output = cfunc(a)
np.testing.assert_allclose(output.asnumpy(), a.asnumpy())
def test_int_mult_3():
mod = Module()
x = var('x', dtype='int32', shape=())
func = Function([x], x * relay.const(3))
cfunc = compile(func, mod)
a = tvm.nd.array(np.array(4, dtype='int32'))
output = cfunc(a)
np.testing.assert_allclose(output.asnumpy(), np.array(12, dtype='int32'))
def test_nat_add():
mod = Module()
p = Prelude(mod)
cfunc = compile(
Function([], p.add(p.s(p.s(p.s(p.z()))), p.s(p.s(p.s(p.s(p.z())))))),
mod)
output = cfunc()
assert nat_to_int(output) == 7
def test_constructor_tag_round_trip():
mod1 = Module()
p1 = Prelude(mod1)
add_nat_definitions(p1)
mod2 = Module()
p2 = Prelude(mod2)
add_nat_definitions(p2)
# ensure hashes match across modules
ctors1 = constructor_list(p1)
ctors2 = constructor_list(p2)
for i in range(len(ctors1)):
tag = ctors1[i].tag
ctor = mod2.get_constructor(tag)
assert ctor == ctors2[i]
assert ctor.name_hint == ctors1[i].name_hint
def test_add_42():
mod = Module()
x = var('x', shape=())
func = Function([x], x + relay.const(42.0))
cfunc = compile(func, mod)
a = tvm.nd.array(np.array(42.0, dtype='float32'))
output = cfunc(a)
np.testing.assert_allclose(output.asnumpy(), np.array(84.0,
dtype='float32'))
def test_map():
mod = Module()
p = Prelude(mod)
f = Var("f")
orig = p.map(f,
p.cons(const(1), p.cons(const(2), p.cons(const(3), p.nil()))))
expected = p.cons(f(const(1)),
p.cons(f(const(2)), p.cons(f(const(3)), p.nil())))
assert alpha_equal(dcpe(orig, mod=mod), expected)
def test_loop():
mod = Module()
t = TypeVar("t")
x = Var("x", t)
loop = GlobalVar("loop")
mod[loop] = Function([x], loop(x), t, [t])
res = dcpe(loop(const(1)), mod=mod)
expected = Call(loop, [const(1)], None, [None])
assert alpha_equal(res, expected)
def test_head_cons():
mod = Module()
p = Prelude(mod)
hd = p.hd
t = TypeVar("t")
x = Var("x", t)
body = hd(p.cons(x, p.nil()))
f = Function([x], body, None, [t])
res = dcpe(f, mod)
assert alpha_equal(res, Function([x], x, t, [t]))
def test_double():
mod = Module()
x = var('x', shape=())
double = GlobalVar('double')
mod[double] = Function([x], x + x)
x = var('x', shape=())
cfunc = compile(Function([x], double(double(x))), mod)
a = tvm.nd.array(np.array(1.5, dtype='float32'))
output = cfunc(a)
np.testing.assert_allclose(output.asnumpy(), np.array(6.0,
dtype='float32'))
def test_tuple():
mod = Module()
cfunc = compile(
Function([],
relay.TupleGetItem(
relay.Tuple([
relay.const(3, dtype='int32'),
relay.const(4.0, dtype='float32')
]), 1)), mod)
np.testing.assert_allclose(cfunc().asnumpy(), np.array(4.0,
dtype='float32'))
def test_global_match_nat_id():
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
nat = p.nat()
x = Var("x", nat)
z_case = Clause(PatternConstructor(p.z, []), p.z())
s_case = Clause(PatternConstructor(p.s, [PatternVar(x)]), p.s(x))
orig = Match(make_nat_expr(p, 3), [z_case, s_case])
res = dcpe(orig, mod=mod)
assert alpha_equal(res, make_nat_expr(p, 3))
def test_add():
mod = Module()
x = var('x', shape=())
y = var('y', shape=())
z = x + y
func = Function([x, y], z)
cfunc = compile(func, mod)
a = tvm.nd.array(np.array(1.0, dtype='float32'))
b = tvm.nd.array(np.array(1.0, dtype='float32'))
c = tvm.nd.array(np.array(2.0, dtype='float32'))
output = cfunc(a, b)
np.testing.assert_allclose(output.asnumpy(), c.asnumpy())
def test_compose():
mod = Module()
p = Prelude(mod)
x = relay.Var('x')
inc = GlobalVar('inc')
mod[inc] = Function([x], p.s(x))
x = relay.Var('x')
func = GlobalVar('func')
f = Function([x], relay.Call(p.compose(inc, p.double), [x]))
mod[func] = f
cfunc = compile(func, mod)
assert nat_to_int(cfunc(p.s(p.s(p.z())))) == 5
def test_recur_sum_global():
mod = Module()
x = var('x', dtype='int32', shape=())
sum = GlobalVar('sum')
c = relay.const(0)
mod[sum] = Function([x],
relay.If(op.less(x, c), c,
x + sum(x - relay.const(1))),
relay.TensorType(dtype='int32', shape=()))
cfunc = compile(Function([], sum(relay.const(10))), mod)
output = cfunc()
np.testing.assert_allclose(output.asnumpy(), np.array(55, dtype='int32'))
def test_recur_sum_local():
mod = Module()
x = var('x', dtype='int32', shape=())
t = relay.TensorType(dtype='int32', shape=())
sum = relay.Var('sum', type_annotation=relay.FuncType([t], t))
c = relay.const(0)
func = Function([x], relay.If(op.less(x, c), c,
x + sum(x - relay.const(1))), t)
body = relay.Let(sum, func, sum(relay.const(10)))
cfunc = compile(Function([], body), mod)
output = cfunc()
np.testing.assert_allclose(output.asnumpy(), np.array(55, dtype='int32'))
def test_swap_loop():
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
nat = p.nat()
x = Var("x", nat)
y = Var("y", nat)
loop = GlobalVar("loop")
mod[loop] = Function([x, y], loop(y, x), nat)
prog = loop(make_nat_expr(p, 1), make_nat_expr(p, 2))
res = dcpe(prog, mod=mod)
assert alpha_equal(prog, res)
def test_nat_id():
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
nat = p.nat()
x = Var("x", nat)
y = Var("y", nat)
nat_id = GlobalVar("nat_id")
mod[nat_id] = Function([x], x)
orig = nat_id(make_nat_expr(p, 3))
res = dcpe(orig, mod=mod)
assert alpha_equal(res, make_nat_expr(p, 3))
def test_loop():
mod = Module()
t = TypeVar("t")
x = Var("x", t)
loop = GlobalVar("loop")
mod[loop] = Function([x], loop(x), t, [t])
expected = Call(loop, [const(1)])
mod["main"] = Function([], expected)
expected = mod["main"].body
call = Function([], loop(const(1)))
res = dcpe(call, mod=mod)
assert alpha_equal(res.body, expected)
def test_mult_op():
mod = Module()
x = var('x', shape=())
y = var('y', shape=())
z = x + y
zz = op.exp(z)
func = Function([x, y], zz)
cfunc = compile(func, mod)
a = tvm.nd.array(np.array(1.0, dtype='float32'))
b = tvm.nd.array(np.array(1.0, dtype='float32'))
output = cfunc(a, b)
np.testing.assert_allclose(output.asnumpy(),
np.exp(a.asnumpy() + b.asnumpy()))
def test_match_nat_id():
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
nat = p.nat()
x = Var("x", nat)
y = Var("y", nat)
nat_id = GlobalVar("nat_id")
z_case = Clause(PatternConstructor(p.z, []), p.z())
s_case = Clause(PatternConstructor(p.s, [PatternVar(y)]), p.s(y))
mod[nat_id] = Function([x], Match(x, [z_case, s_case]))
orig = nat_id(make_nat_expr(p, 3))
res = dcpe(orig, mod=mod)
assert alpha_equal(res, make_nat_expr(p, 3))
def test_map():
mod = Module()
p = Prelude(mod)
f = GlobalVar("f")
t = TypeVar("t")
a = Var("a", t)
mod[f] = Function([a], a, t, [t])
orig = p.map(f, p.cons(const(1), p.cons(const(2), p.cons(const(3), p.nil()))))
expected = p.cons((const(1)), p.cons((const(2)), p.cons((const(3)), p.nil())))
expected = Function([], expected)
mod["main"] = expected
expected = mod["main"]
orig = Function([], orig)
res = dcpe(orig, mod=mod)
assert alpha_equal(res.body, expected.body)
def test_abs():
mod = Module()
x = var('x', shape=())
func = Function([x],
relay.If(op.less(x, relay.const(0.0)),
relay.const(-1.0) * x, x))
cfunc = compile(func, mod)
a = tvm.nd.array(np.array(12.0, dtype='float32'))
output = cfunc(a)
np.testing.assert_allclose(output.asnumpy(), np.array(12.0,
dtype='float32'))
a = tvm.nd.array(np.array(-34.0, dtype='float32'))
output = cfunc(a)
np.testing.assert_allclose(output.asnumpy(), np.array(34.0,
dtype='float32'))
def double_example():
# Declare a Relay module.
mod = Module()
# Implement the double function.
x = var('x', shape=())
double = GlobalVar('double')
mod[double] = Function([x], x + x)
# Generate a function which calls double twice.
x = var('x', shape=())
f = Function([x], double(double(x)))
# Compile the function.
cfunc = compile(f, mod)
a = tvm.nd.array(np.array(1.5, dtype='float32'))
print(cfunc(a).asnumpy())
def test_constructor_tag_differences():
# ensure that if we have the type data for a given ADT, the tags
# for the constructors of the *same ADT* are simple offsets from
# each other
mod = Module()
p = Prelude(mod)
add_nat_definitions(p)
adts = adt_list(p)
for adt in adts:
data = mod[adt]
for i in range(len(data.constructors) - 1):
ctor1 = data.constructors[i]
ctor2 = data.constructors[i + 1]
assert ctor2.tag - ctor1.tag == 1
# make sure there is something present at the MSB
assert ctor1.tag - i != 0
assert ctor2.tag - (i + 1) != 0
def test_local_local_rec_outer_scope():
mod = Module()
x = var('x', dtype='int32', shape=())
t = relay.TensorType(dtype='int32', shape=())
sum = relay.Var('sum', type_annotation=relay.FuncType([t], t))
c = relay.const(0)
# we define a locally recursive function inside another function's scope
# and have that function return the closure of the locally recursive function
inner_func = Function([x],
relay.If(op.less(x, c), c,
x + sum(x - relay.const(1))), t)
outer_func_body = relay.Let(sum, inner_func, sum)
outer_func = Function([], outer_func_body)
f = relay.Var('f')
body = relay.Let(f, outer_func(), f(relay.const(10)))
cfunc = compile(Function([], body), mod)
output = cfunc()
np.testing.assert_allclose(output.asnumpy(), np.array(55, dtype='int32'))
