def test_dual_op():
"""Program:
fn (x : Tensor[f32, (10, 10)]) {
let t1 = log(x);
let t2 = add(t1, x);
return t1;
}
"""
b = IRBuilder()
with b.function(('x', tensor_type(10, 10))) as func:
x, = func.param_ids()
t1 = b.let('t1', log(x))
t2 = b.let('t2', add(t1, x))
b.ret(t2)
assert_has_type(func.to_func(), func_type(['float32'], 'float32'))
def check_single_op(opfunc):
"Program: fn (x : float32) { let t1 = f(x); t1 }"
b = IRBuilder()
with b.function(('x', 'float32')) as func:
x, = func.param_ids()
t1 = b.let('t1', opfunc(x))
b.ret(t1)
assert_has_type(func.to_func(), func_type(['float32'], 'float32'))
def test_monomorphic_let():
"Program: let x = 1; return x"
b = IRBuilder()
x = b.let('x', 1.0, value_type=scalar_type('float64'))
b.ret(x)
prog, env = b.get()
assert_has_type(prog, scalar_type('float64'))
def test_let():
b = IRBuilder()
x = b.let('x', 1)
b.ret(x)
prog, _ = b.get()
assert isinstance(prog, Let)
var = prog.var
value = prog.value
assert var.name_hint == 'x'
assert var == prog.body
assert isinstance(value, Constant)
assert value.data.asnumpy() == np.array(1)
def test_decl():
"""Program:
def f(x : Tensor[f32, (10, 10)]) {
let lx = log(x);
return lx;
}
"""
b = IRBuilder()
x = b.param('x')
with b.decl('f', x):
lx = b.let('lx', log(x))
b.ret(lx)
_, env = b.get()
assert_decl_has_type(env, 'f', func_type(['float32'], 'float32'))