def test_raw_tensor(): from megengine.core.ops.builtin import Elemwise x = np.random.rand(10).astype("float32") xx = Tensor(x) (yy, ) = apply(Elemwise(Elemwise.Mode.MUL), xx, xx) np.testing.assert_allclose(x * x, yy.numpy()) (yy, ) = apply(Elemwise(Elemwise.Mode.MUL), xx, xx) np.testing.assert_allclose(x * x, yy.numpy())
def _elwise(mode): op = Elemwise(mode) def f(*args): (result, ) = apply(op, *args) return result return f
def test_elemwise_relu_backward_fn(): op = Elemwise(Elemwise.Mode.RELU) attr = TensorAttr() attr.dtype = "float32" attr.comp_node = "xpux" result = imperative.make_backward_graph(op, [attr], [True], [True]) backward_graph, save_for_backward_mask, input_has_grad = result assert save_for_backward_mask == [False, True, True], save_for_backward_mask
def test_replace_vars(): g = mgb_graph.Graph() g.options.async_exec_level = 0b100 device = "xpux" dtype = np.float32 a = mgb_graph.InputNode(device=device, dtype=dtype, graph=g) const = g.make_const(1.234, device=device) add_op = Elemwise(Elemwise.Mode.ADD) mul_op = Elemwise(Elemwise.Mode.MUL) a_plus_a = apply_normal_varnode(add_op, a.outputs[0], a.outputs[0])[0] a_plus_a_mul_const = apply_normal_varnode(mul_op, a_plus_a, const)[0] rst = apply_normal_varnode(add_op, a_plus_a_mul_const, a.outputs[0])[0] (new, ) = cgtools.replace_vars([rst._node], {const._node: a_plus_a._node}) out = mgb_graph.OutputNode(mgb_graph.VarNode(new)) func = g.compile(out.outputs[0]) func.execute() x = make_dev_tensor(5.0, device=device) a.set_value(x) res = out.get_value().numpy() np.testing.assert_equal(res, np.array([105.0]))
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())
def test_exception(): err_msg = "QwQ" def throw_exc(): raise RuntimeError(err_msg) g = mgb_graph.Graph() x, _ = mgb_graph.input_callback(throw_exc, device="xpux", dtype="float32", graph=g) neg = Elemwise(Elemwise.Mode.NEGATE) y = mgb_graph.OutputNode(mgb_graph.apply_normal_varnode(neg, x)[0]) f = g.compile(y.outputs[0]) try: f.execute() y.get_value() except Exception as exc: assert err_msg in str(exc)
def test_opr_attr(): from megengine.core.ops.builtin import Elemwise assert Elemwise(Elemwise.Mode.ADD) == Elemwise(Elemwise.Mode.ADD)
def elemwise(*args, mode): from megengine.core.ops.builtin import Elemwise return apply(Elemwise(mode), *args)
def elemwise(*args, mode): from megengine.core._imperative_rt.imperative import apply_op from megengine.core.ops.builtin import Elemwise return apply_op(Elemwise(mode), args)
def test_opr_attr(): from megengine.core.ops.builtin import Elemwise assert Elemwise(mode="add") == Elemwise(mode="add")
def forward(self, x, y): op = Elemwise("ADD") return apply(op, x, y)[0]