def test_straightforward(self):
x, y, z = floats("xyz")
e = mul(add(x, y), true_div(x, y))
C = Composite([x, y], [e])
c = C.make_node(x, y)
# print c.c_code(['x', 'y'], ['z'], dict(id = 0))
g = FunctionGraph([x, y], [c.out])
fn = DualLinker().accept(g).make_function()
assert fn(1.0, 2.0) == 1.5
def test_with_constants(self):
x, y, z = floats("xyz")
e = mul(add(70.0, y), true_div(x, y))
C = Composite([x, y], [e])
c = C.make_node(x, y)
assert "70.0" in c.op.c_code(c, "dummy", ["x", "y"], ["z"], dict(id=0))
# print c.c_code(['x', 'y'], ['z'], dict(id = 0))
g = FunctionGraph([x, y], [c.out])
fn = DualLinker().accept(g).make_function()
assert fn(1.0, 2.0) == 36.0
def test_many_outputs(self):
x, y, z = floats("xyz")
e0 = x + y + z
e1 = x + y * z
e2 = x / y
C = Composite([x, y, z], [e0, e1, e2])
c = C.make_node(x, y, z)
# print c.c_code(['x', 'y', 'z'], ['out0', 'out1', 'out2'], dict(id = 0))
g = FunctionGraph([x, y, z], c.outputs)
fn = DualLinker().accept(g).make_function()
assert fn(1.0, 2.0, 3.0) == [6.0, 7.0, 0.5]
def test_flatten(self):
# Test that we flatten multiple Composite.
x, y, z = floats("xyz")
C = Composite([x, y], [x + y])
CC = Composite([x, y], [C(x * y, y)])
assert not isinstance(CC.outputs[0].owner.op, Composite)
# Test with multiple outputs
CC = Composite([x, y, z], [C(x * y, y), C(x * z, y)])
# We don't flatten that case.
assert isinstance(CC.outputs[0].owner.op, Composite)
def test_numba_Composite(inputs, input_values):
x_s = aes.float64("x")
y_s = aes.float64("y")
comp_op = Elemwise(
Composite([x_s, y_s], [x_s + y_s * 2 + aes.exp(x_s - y_s)]))
out_fg = FunctionGraph(inputs, [comp_op(*inputs)])
compare_numba_and_py(out_fg, input_values)
def test_composite_clone_float32(self):
def has_f16(comp):
if any(v.type == float16 for v in comp.fgraph.variables):
return True
return False
w = int8()
x = float16()
y = float32()
cz = Composite([x, y], [tanh(x + cast(y, "float16"))])
c = Composite(
[w, x, y],
[
cz(x, y) - cz(x, y)**2 + cast(x, "int16") +
cast(x, "float32") + cast(w, "float16") -
constant(np.float16(1.0))
],
)
assert has_f16(c)
nc = c.clone_float32()
assert not has_f16(nc)
v = uint8()
w = float16()
x = float16()
y = float16()
z = float16()
c = Composite([v, w, x, y, z], [switch(v, mul(w, x, y), z)])
assert has_f16(c)
nc = c.clone_float32()
assert not has_f16(nc)
def test_composite_printing(self):
x, y, z = floats("xyz")
e0 = x + y + z
e1 = x + y * z
e2 = x / y
e3 = x // 5
e4 = -x
e5 = x - y
e6 = x**y + (-z)
e7 = x % 3
C = Composite([x, y, z], [e0, e1, e2, e3, e4, e5, e6, e7])
c = C.make_node(x, y, z)
g = FunctionGraph([x, y, z], c.outputs)
DualLinker().accept(g).make_function()
assert str(g) == ("FunctionGraph(*1 -> Composite{((i0 + i1) + i2),"
" (i0 + (i1 * i2)), (i0 / i1), "
"(i0 // ScalarConstant{5}), "
"(-i0), (i0 - i1), ((i0 ** i1) + (-i2)),"
" (i0 % ScalarConstant{3})}(x, y, z), "
"*1::1, *1::2, *1::3, *1::4, *1::5, *1::6, *1::7)")
def test_jax_Composite(x, y, x_val, y_val):
x_s = aes.float64("x")
y_s = aes.float64("y")
comp_op = Elemwise(Composite([x_s, y_s], [x_s + y_s * 2 + aes.exp(x_s - y_s)]))
out = comp_op(x, y)
out_fg = FunctionGraph([x, y], [out])
test_input_vals = [
x_val.astype(config.floatX),
y_val.astype(config.floatX),
]
_ = compare_jax_and_py(out_fg, test_input_vals)
