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]