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 f(x):
return Tensor(mgb.opr.assert_equal(x._symvar, x._symvar + 1))
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())