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_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 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'))
def _less(children, attrs, odtype='float32'):
out_type = attrs.get_str('out_type', None)
if out_type:
return op.less(children[0], children[1]).astype(out_type)
else:
return op.less(children[0], children[1])
