def test_io():
g = mgb_graph.Graph()
x = Tensor(np.random.randn(3).astype("float32"), device="xpux")._dev_tensor()
vx, _ = mgb_graph.input_callback(
lambda: x, device=x.comp_node, dtype=x.dtype, graph=g
)
y = Future()
v = mgb_graph.output_callback(y.set_result, vx)
f = g.compile(v)
f()
np.testing.assert_equal(x.numpy(), y.result().numpy())
def test_op():
g = mgb_graph.Graph()
x = Tensor(np.random.randn(10).astype("float32"), device="xpux")._dev_tensor()
v, _ = mgb_graph.input_callback(
lambda: x, device=x.comp_node, dtype=x.dtype, graph=g
)
neg = Elemwise(Elemwise.Mode.NEGATE)
v = mgb_graph.apply_normal_varnode(neg, v)[0]
y = Future()
v = mgb_graph.output_callback(y.set_result, v)
f = g.compile(v)
f()
np.testing.assert_equal(x.numpy(), -y.result().numpy())