def test_add_output(): a = Tensor([1.0, 2.0]) b = Tensor([3.0, 4.0]) @trace(symbolic=True, capture_as_const=True) def fwd(a, b): return (a + b) * 2 fwd(a, b) orig_model = io.BytesIO() fwd.dump( orig_model, arg_names=["a", "b"], output_names="o", optimize_for_inference=False ) orig_model.seek(0) net = Net.load(orig_model) var_a = net.var_filter.name("a").as_unique() var_b = net.var_filter.name("b").as_unique() y = F.add(var_a, var_b) y = F.sigmoid(y) y.name = "o1" net.add_output(y) modified_model = io.BytesIO() net.dump(modified_model) modified_model.seek(0) g = GraphInference(modified_model) out = g.run(a.numpy(), b.numpy()) np.testing.assert_equal(out["o"], ((a + b) * 2).numpy()) np.testing.assert_equal(out["o1"], (F.sigmoid((a + b))).numpy())
def test_set_subtensor(): x = Tensor([1, 2, 3]) x[:] = [1, 1, 1] np.testing.assert_almost_equal(x.numpy(), [1, 1, 1], decimal=6) x[[0, 2]] = [3, 2] np.testing.assert_almost_equal(x.numpy(), [3, 1, 2], decimal=6) x[1:3] = [4, 5] np.testing.assert_almost_equal(x.numpy(), [3, 4, 5], decimal=6)
def _shared_nd_test(bit, low_bit_type): max_value = (1 << bit) - 1 min_value = 1 - (1 << bit) data = np.arange(min_value, max_value + 2, 2, dtype=low_bit_type) snd = Tensor(data, dtype=low_bit_type, device="xpux") np.testing.assert_allclose(snd.numpy(), range(min_value, max_value + 2, 2)) data = np.arange(min_value, max_value + 2, 4, dtype=low_bit_type) snd = Tensor(data, dtype=low_bit_type, device="xpux") np.testing.assert_allclose(snd.numpy(), range(min_value, max_value + 2, 4))
def test_shared_nd(): data = np.array([-3.4, 1.394683, 2.323497, -7.439948, -5.2397], dtype=bfloat16) snd = Tensor(data, dtype=bfloat16, device="xpux") assert snd.numpy().dtype == bfloat16 np.testing.assert_allclose(snd.numpy(), [-3.40625, 1.398438, 2.328125, -7.4375, -5.25], atol=1e-6) data = np.array([-9.34964, -8.342, 9.4385, 0.18746, 1.48], dtype=bfloat16) snd = Tensor(data, dtype=bfloat16, device="xpux") np.testing.assert_allclose(snd.numpy(), [-9.375, -8.3125, 9.4375, 0.1875, 1.476562], atol=1e-6)
def test_serialization(): x = Tensor([1, 2, 3], dtype=np.float32) newargs = x.__getnewargs__() states = x.__getstate__() assert np.all(newargs[0] == x.numpy()) assert newargs[1] == x.dtype assert newargs[2] == x.device.logical_name assert not states x.qparams states = x.__getstate__() assert len(states.keys()) == 1 assert states["qparams"] == x.qparams
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_add_remove_output(): a = Tensor([1.0, 2.0]) b = Tensor([3.0, 4.0]) @trace(symbolic=True, capture_as_const=True) def fwd(a, b): return (a + b) * 2, (a - b) fwd(a, b) orig_model = io.BytesIO() fwd.dump( orig_model, arg_names=["a", "b"], output_names=["o1", "o2"], optimize_for_inference=False, ) orig_model.seek(0) net = Net.load(orig_model) var_a = net.var_filter.name("a").as_unique() var_b = net.var_filter.name("b").as_unique() y1 = (var_a + var_b) * 3 y2 = F.sigmoid(var_a + var_b) net.remove_output(*net.output_vars) y1.name = "new_o1" y2.name = "new_o2" net.add_output(y1, y2) modified_model = io.BytesIO() net.dump(modified_model) modified_model.seek(0) g = GraphInference(modified_model) out = g.run(a.numpy(), b.numpy()) np.testing.assert_equal(out["new_o1"], ((a + b) * 3).numpy()) np.testing.assert_almost_equal(out["new_o2"], (F.sigmoid((a + b))).numpy())
def test_dump_cond_take(): a = Tensor([1.0, 2.0]) @trace(symbolic=True, capture_as_const=True) def fwd(a): return F.cond_take(a > 1, a) fwd(a) orig_model = io.BytesIO() fwd.dump( orig_model, arg_names=["a"], output_names=["o1", "o2"], optimize_for_inference=False, ) orig_model.seek(0) net = Net.load(orig_model) var_a = net.input_vars[0] val, idx = F.cond_take(var_a > 1, var_a) net.remove_output(*net.output_vars) val.name = "value" idx.name = "index" net.add_output(val, idx) modified_model = io.BytesIO() net.dump(modified_model) modified_model.seek(0) g = GraphInference(modified_model) out = g.run(a.numpy()) data = a.numpy() mask = a.numpy() > 1 np.testing.assert_equal(out["index"], np.where(mask.reshape(-1))[0]) np.testing.assert_equal(out["value"], data[mask])
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_io2(): g = mgb_graph.Graph() g.options.async_exec_level = 0b100 dtype, device = "float32", "xpux" px = mgb_graph.InputNode(device=device, dtype=dtype, graph=g) py = mgb_graph.OutputNode(px.outputs[0]) f = g.compile(py.outputs[0]) for _ in range(3): f.execute() x = Tensor(np.random.randn(10).astype(dtype), device=device)._dev_tensor() px.set_value(x) y = py.get_value() np.testing.assert_equal(x.numpy(), y.numpy()) f.wait()
def test_qint_new_tensor(dtype, dtype_name): convert_to_dtype = eval("convert_to_%s" % dtype_name) convert_from_dtype = eval("convert_from_%s" % dtype_name) shape = (3, 3, 3) data = np.random.random(shape).astype(np.float32) * 5 - 1 # create a new Tensor with quint8 dtype inp = Tensor(convert_to_dtype(data, dtype), dtype=dtype) _check_result_attr(inp, dtype, dtype_name, dtype_name.startswith("qu")) np.testing.assert_equal(inp.numpy(), convert_to_dtype(data, dtype)) # convert from quint8 to float32 inp_float = inp.astype("float32") assert inp_float.dtype == np.float32 np.testing.assert_equal(inp_float.numpy(), convert_from_dtype(convert_to_dtype(data, dtype)))
def test_set_value(): v0 = np.random.random((2, 3)).astype(np.float32) param = Tensor(v0) v1 = np.random.random((2, 3)).astype(np.float32) param[...] = v1 np.testing.assert_allclose(param.numpy(), v1, atol=5e-6)
def test_matmul(): A = Tensor(np.random.rand(5, 7).astype("float32")) B = Tensor(np.random.rand(7, 10).astype("float32")) C = A @ B np.testing.assert_almost_equal(C.numpy(), A.numpy() @ B.numpy(), decimal=6)
def test_advance_indexing_with_bool(): a = np.arange(9).reshape(3, 3).astype(np.float32) b = np.array([1, 2, 3]) c = np.array([1, 2, 3]) aa = Tensor(a) bb = Tensor(b) cc = Tensor(c) np.testing.assert_equal(a[b == 1, c == 2], aa[bb == 1, cc == 2].numpy()) a[b == 1, c == 2] = -1.0 aa[bb == 1, cc == 2] = -1.0 np.testing.assert_equal(a, aa.numpy()) a = np.arange(9).reshape(3, 3).astype(np.float32) b = np.array([False, True, True]) c = np.array([2, 0]).astype(np.int32) aa = Tensor(a) bb = Tensor(b) cc = Tensor(c) np.testing.assert_equal(a[b, c], aa[bb, cc].numpy()) a[b, c] = -1.0 aa[bb, cc] = -1.0 np.testing.assert_equal(a, aa.numpy()) d = np.array([-1, -2], dtype=np.float32) dd = Tensor(d) a[b, c] = d aa[bb, cc] = dd np.testing.assert_equal(a, aa.numpy()) a = np.ones((2, 2)) b = np.array([[True, False], [False, True]]) aa = Tensor(a) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy()) b[:] = True bb[:] = True np.testing.assert_equal(a[b], aa[bb].numpy()) np.testing.assert_equal(a[:, [True, False]], aa[:, [True, False]].numpy()) a = np.array([[True, False], [False, True]]) b = np.array([1]) aa = Tensor(a) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy()) b = np.array([[True, True], [False, True]]) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy()) a[b] = False aa[bb] = False np.testing.assert_equal(a, aa.numpy()) # XXX: trace does not expect empty condtake tensor if not use_symbolic_shape(): a = np.ones((2, 2), dtype=np.int32) b = np.array([[False, False], [False, False]]) aa = Tensor(a) bb = Tensor(b) np.testing.assert_equal(a[b], aa[b].numpy()) np.testing.assert_equal(a[b], aa[bb].numpy()) b = np.array([False, False]) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy().reshape(a[b].shape)) # FIXME a = np.arange(576).reshape(2, 3, 4, 3, 4, 2).astype("int32") aa = Tensor(a) b = (np.random.sample((2, 3, 4)) > 0.5).astype("bool") bb = Tensor(b) np.testing.assert_equal(a[b, :, 0:4:2], aa[bb, :, 0:4:2].numpy()) b = (np.random.sample((4, 3, 4)) > 0.5).astype("bool") bb = Tensor(b) np.testing.assert_equal(a[..., b, 0:2], aa[..., bb, 0:2].numpy()) b = (np.random.sample((3, 4, 3)) > 0.5).astype("bool") bb = Tensor(b) np.testing.assert_equal( a[:, b, 0:2, [True, False]], aa[:, bb, 0:2, [True, False]].numpy() )
def test_advance_indexing_high_level(): x = np.arange(25).reshape(5, 5).astype("int32") d = np.arange(15).reshape(3, 5).astype("int32") xx = Tensor(x) np.testing.assert_equal(x[1, :], xx[1, :].numpy()) np.testing.assert_equal(x[:, 1], xx[:, 1].numpy()) np.testing.assert_equal(x[1:3, :], xx[1:3, :].numpy()) np.testing.assert_equal(x[:, :], xx[:, :].numpy()) np.testing.assert_equal(x[1, 1], xx[1, 1].numpy()) yy = xx[(0, 4, 2), :] np.testing.assert_equal(x[(0, 4, 2), :], yy.numpy()) x_ = x.copy() x_[(0, 4, 2), :] = d xx_ = Tensor(xx) xx_[(0, 4, 2), :] = d np.testing.assert_equal(x_, xx_.numpy()) x = np.arange(27).reshape(3, 3, 3).astype("int32") xx = Tensor(x) np.testing.assert_equal(x[1, :, :], xx[1, :, :].numpy()) np.testing.assert_equal(x[1, :, 1], xx[1, :, 1].numpy()) np.testing.assert_equal(x[1, 0:1, :], xx[1, 0:1, :].numpy()) np.testing.assert_equal(x[0:1, 1, 1], xx[0:1, 1, 1].numpy()) np.testing.assert_equal(x[:, 1, 1], xx[:, 1, 1].numpy()) np.testing.assert_equal(x[:, 1], xx[:, 1].numpy()) np.testing.assert_equal(x[1, 1:2], xx[1, 1:2].numpy()) x_ = x.copy() x_[1, 1, 1] = -1 xx[1, 1, 1] = -1 np.testing.assert_equal(x_, xx.numpy()) x_[:, 1, 1] = -2 xx[:, 1, 1] = x_[:, 1, 1] np.testing.assert_equal(x_, xx.numpy()) x_[0:1, :, 1] = -3 xx[0:1, :, 1] = x_[0:1, :, 1] np.testing.assert_equal(x_, xx.numpy()) x_[0:1, :, 1] = -4 y = Tensor(x_) xx[0:1, :, 1] = y[0:1, :, 1] np.testing.assert_equal(y.numpy(), xx.numpy()) x[:] = 1 xx[:] = 1 np.testing.assert_equal(x, xx.numpy()) x = np.arange(9).reshape(3, 3).astype("int32") xx = Tensor(x) y = np.array([1, 2]) yy = Tensor(y) np.testing.assert_equal(x[:, y[0]], xx[:, y[0]].numpy()) np.testing.assert_equal(x[:, y[0]], xx[:, yy[0]].numpy()) np.testing.assert_equal(x[:, y], xx[:, y].numpy()) np.testing.assert_equal(x[:, y], xx[:, yy].numpy()) x_ = x.copy() x_[:, y[0]] = -1 xx_ = Tensor(x_) xx[:, yy[0]] = xx_[:, yy[0]] np.testing.assert_equal(x_, xx.numpy()) x_[:, y] = -1 xx_ = Tensor(x_) xx[:, yy] = xx_[:, yy] np.testing.assert_equal(x_, xx.numpy()) x = np.arange(9).reshape(3, 3).astype("int32") xx = Tensor(x) y = np.array([1]) yy = Tensor(y) np.testing.assert_equal(x[:, y[0]], xx[:, y[0]].numpy()) np.testing.assert_equal(x[:, y[0]], xx[:, yy[0]].numpy()) np.testing.assert_equal(x[:, y], xx[:, y].numpy()) np.testing.assert_equal(x[:, y], xx[:, yy].numpy()) x = np.arange(9).reshape(3, 3).astype("int32") xx = Tensor(x) np.testing.assert_equal(x[[0, 1], 0], xx[[0, 1], 0].numpy()) np.testing.assert_equal(x[0:2, 0], xx[0:2, 0].numpy())
def test_advance_indexing_with_bool(test_varnode): if test_varnode: network = Network() else: network = None a = np.arange(9).reshape(3, 3).astype(np.float32) b = np.array([1, 2, 3]) c = np.array([1, 2, 3]) aa = make_tensor(a, network) bb = make_tensor(b, network) cc = make_tensor(c, network) np.testing.assert_equal(a[b == 1, c == 2], get_value(aa[bb == 1, cc == 2])) a[b == 1, c == 2] = -1.0 aa[bb == 1, cc == 2] = -1.0 np.testing.assert_equal(a, get_value(aa)) a = np.arange(9).reshape(3, 3).astype(np.float32) b = np.array([False, True, True]) c = np.array([2, 0]).astype(np.int32) aa = Tensor(a) bb = Tensor(b) cc = Tensor(c) np.testing.assert_equal(a[b, c], aa[bb, cc].numpy()) a[b, c] = -1.0 aa[bb, cc] = -1.0 np.testing.assert_equal(a, aa.numpy()) d = np.array([-1, -2], dtype=np.float32) dd = Tensor(d) a[b, c] = d aa[bb, cc] = dd np.testing.assert_equal(a, aa.numpy()) a = np.ones((2, 2)) b = np.array([[True, False], [False, True]]) aa = Tensor(a) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy()) b[:] = True bb[:] = True np.testing.assert_equal(a[b], aa[bb].numpy()) np.testing.assert_equal(a[:, [True, False]], aa[:, [True, False]].numpy()) a = np.array([[True, False], [False, True]]) b = np.array([1]) aa = Tensor(a) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy()) b = np.array([[True, True], [False, True]]) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy()) a[b] = False aa[bb] = False np.testing.assert_equal(a, aa.numpy()) a = np.ones((2, 2), dtype=np.int32) b = np.array([[False, False], [False, False]]) aa = Tensor(a) bb = Tensor(b) np.testing.assert_equal(a[b], aa[b].numpy()) np.testing.assert_equal(a[b], aa[bb].numpy()) b = np.array([False, False]) bb = Tensor(b) np.testing.assert_equal(a[b], aa[bb].numpy().reshape(a[b].shape)) a = np.arange(576).reshape(2, 3, 4, 3, 4, 2).astype("int32") aa = Tensor(a) b = (np.random.sample((2, 3, 4)) > 0.5).astype("bool") bb = Tensor(b) np.testing.assert_equal(a[b, :, 0:4:2], aa[bb, :, 0:4:2].numpy()) b = (np.random.sample((4, 3, 4)) > 0.5).astype("bool") bb = Tensor(b) np.testing.assert_equal(a[..., b, 0:2], aa[..., bb, 0:2].numpy()) b = (np.random.sample((3, 4, 3)) > 0.5).astype("bool") bb = Tensor(b) np.testing.assert_equal(a[:, b, 0:2, [True, False]], aa[:, bb, 0:2, [True, False]].numpy())