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())
