def test_squeeze(self):
torch.manual_seed(1)
data = ht.random.randn(1, 4, 5, 1, device=ht_device)
# 4D local tensor, no axis
result = ht.squeeze(data)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.float64)
self.assertEqual(result._DNDarray__array.dtype, torch.float64)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.lshape, (4, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze()).all())
# 4D local tensor, major axis
result = ht.squeeze(data, axis=0)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.float64)
self.assertEqual(result._DNDarray__array.dtype, torch.float64)
self.assertEqual(result.shape, (4, 5, 1))
self.assertEqual(result.lshape, (4, 5, 1))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze(0)).all())
# 4D local tensor, minor axis
result = ht.squeeze(data, axis=-1)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.float64)
self.assertEqual(result._DNDarray__array.dtype, torch.float64)
self.assertEqual(result.shape, (1, 4, 5))
self.assertEqual(result.lshape, (1, 4, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze(-1)).all())
# 4D local tensor, tuple axis
result = data.squeeze(axis=(0, -1))
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.float64)
self.assertEqual(result._DNDarray__array.dtype, torch.float64)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.lshape, (4, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze()).all())
# 4D split tensor, along the axis
# TODO: reinstate this test of uneven dimensions distribution
# after update to Allgatherv implementation (Issue #273 depending on #233)
# data = ht.array(ht.random.randn(1, 4, 5, 1), split=1)
# result = ht.squeeze(data, axis=-1)
# self.assertIsInstance(result, ht.DNDarray)
# # TODO: the following works locally but not when distributed,
# #self.assertEqual(result.dtype, ht.float32)
# #self.assertEqual(result._DNDarray__array.dtype, torch.float32)
# self.assertEqual(result.shape, (1, 12, 5))
# self.assertEqual(result.lshape, (1, 12, 5))
# self.assertEqual(result.split, 1)
# 3D split tensor, across the axis
size = ht.MPI_WORLD.size * 2
data = ht.triu(ht.ones((1, size, size), split=1, device=ht_device), k=1)
result = ht.squeeze(data, axis=0)
self.assertIsInstance(result, ht.DNDarray)
# TODO: the following works locally but not when distributed,
# self.assertEqual(result.dtype, ht.float32)
# self.assertEqual(result._DNDarray__array.dtype, torch.float32)
self.assertEqual(result.shape, (size, size))
self.assertEqual(result.lshape, (size, size))
# self.assertEqual(result.split, None)
# check exceptions
with self.assertRaises(ValueError):
data.squeeze(axis=(0, 1))
with self.assertRaises(TypeError):
data.squeeze(axis=1.1)
with self.assertRaises(TypeError):
data.squeeze(axis="y")
with self.assertRaises(ValueError):
ht.argmin(data, axis=-4)
def test_triu(self):
local_ones = ht.ones((5, ))
# 1D case, no offset, data is not split, module-level call
result = ht.triu(local_ones)
comparison = torch.ones((5, 5), device=self.device.torch_device).triu()
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (5, 5))
self.assertEqual(result.lshape, (5, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == comparison).all())
# 1D case, positive offset, data is not split, module-level call
result = ht.triu(local_ones, k=2)
comparison = torch.ones(
(5, 5), device=self.device.torch_device).triu(diagonal=2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (5, 5))
self.assertEqual(result.lshape, (5, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == comparison).all())
# 1D case, negative offset, data is not split, module-level call
result = ht.triu(local_ones, k=-2)
comparison = torch.ones(
(5, 5), device=self.device.torch_device).triu(diagonal=-2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (5, 5))
self.assertEqual(result.lshape, (5, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == comparison).all())
local_ones = ht.ones((4, 5))
# 2D case, no offset, data is not split, method
result = local_ones.triu()
comparison = torch.ones((4, 5), device=self.device.torch_device).triu()
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.lshape, (4, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == comparison).all())
# 2D case, positive offset, data is not split, method
result = local_ones.triu(k=2)
comparison = torch.ones(
(4, 5), device=self.device.torch_device).triu(diagonal=2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.lshape, (4, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == comparison).all())
# 2D case, negative offset, data is not split, method
result = local_ones.triu(k=-2)
comparison = torch.ones(
(4, 5), device=self.device.torch_device).triu(diagonal=-2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.lshape, (4, 5))
self.assertEqual(result.split, None)
self.assertTrue((result._DNDarray__array == comparison).all())
local_ones = ht.ones((3, 4, 5, 6))
# 2D+ case, no offset, data is not split, module-level call
result = local_ones.triu()
comparison = torch.ones((5, 6), device=self.device.torch_device).triu()
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (3, 4, 5, 6))
self.assertEqual(result.lshape, (3, 4, 5, 6))
self.assertEqual(result.split, None)
for i in range(3):
for j in range(4):
self.assertTrue(
(result._DNDarray__array[i, j] == comparison).all())
# 2D+ case, positive offset, data is not split, module-level call
result = local_ones.triu(k=2)
comparison = torch.ones(
(5, 6), device=self.device.torch_device).triu(diagonal=2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (3, 4, 5, 6))
self.assertEqual(result.lshape, (3, 4, 5, 6))
self.assertEqual(result.split, None)
for i in range(3):
for j in range(4):
self.assertTrue(
(result._DNDarray__array[i, j] == comparison).all())
# # 2D+ case, negative offset, data is not split, module-level call
result = local_ones.triu(k=-2)
comparison = torch.ones(
(5, 6), device=self.device.torch_device).triu(diagonal=-2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (3, 4, 5, 6))
self.assertEqual(result.lshape, (3, 4, 5, 6))
self.assertEqual(result.split, None)
for i in range(3):
for j in range(4):
self.assertTrue(
(result._DNDarray__array[i, j] == comparison).all())
distributed_ones = ht.ones((5, ), split=0)
# 1D case, no offset, data is split, method
result = distributed_ones.triu()
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (5, 5))
self.assertEqual(result.split, 1)
self.assertEqual(result.lshape[0], 5)
self.assertLessEqual(result.lshape[1], 5)
self.assertTrue(result.sum(), 15)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
# 1D case, positive offset, data is split, method
result = distributed_ones.triu(k=2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (5, 5))
self.assertEqual(result.split, 1)
self.assertEqual(result.lshape[0], 5)
self.assertLessEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 6)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
# 1D case, negative offset, data is split, method
result = distributed_ones.triu(k=-2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (5, 5))
self.assertEqual(result.split, 1)
self.assertEqual(result.lshape[0], 5)
self.assertLessEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 22)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
distributed_ones = ht.ones((4, 5), split=0)
# 2D case, no offset, data is horizontally split, method
result = distributed_ones.triu()
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.split, 0)
self.assertLessEqual(result.lshape[0], 4)
self.assertEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 14)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
# # 2D case, positive offset, data is horizontally split, method
result = distributed_ones.triu(k=2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.split, 0)
self.assertLessEqual(result.lshape[0], 4)
self.assertEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 6)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
# # 2D case, negative offset, data is horizontally split, method
result = distributed_ones.triu(k=-2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.split, 0)
self.assertLessEqual(result.lshape[0], 4)
self.assertEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 19)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
distributed_ones = ht.ones((4, 5), split=1)
# 2D case, no offset, data is vertically split, method
result = distributed_ones.triu()
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.split, 1)
self.assertEqual(result.lshape[0], 4)
self.assertLessEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 14)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
# 2D case, positive offset, data is horizontally split, method
result = distributed_ones.triu(k=2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.split, 1)
self.assertEqual(result.lshape[0], 4)
self.assertLessEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 6)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
# 2D case, negative offset, data is horizontally split, method
result = distributed_ones.triu(k=-2)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.split, 1)
self.assertEqual(result.lshape[0], 4)
self.assertLessEqual(result.lshape[1], 5)
self.assertEqual(result.sum(), 19)
if result.comm.rank == 0:
self.assertTrue(result._DNDarray__array[-1, 0] == 0)
if result.comm.rank == result.shape[0] - 1:
self.assertTrue(result._DNDarray__array[0, -1] == 1)
def test_argmin(self):
torch.manual_seed(1)
data = ht.random.randn(3, 4, 5)
# 3D local tensor, no axis
result = ht.argmin(data)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.int64)
self.assertEqual(result._DNDarray__array.dtype, torch.int64)
self.assertEqual(result.shape, (1, ))
self.assertEqual(result.lshape, (1, ))
self.assertEqual(result.split, None)
self.assertTrue(
(result._DNDarray__array == data._DNDarray__array.argmin()).all())
# 3D local tensor, major axis
result = ht.argmin(data, axis=0)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.int64)
self.assertEqual(result._DNDarray__array.dtype, torch.int64)
self.assertEqual(result.shape, (4, 5))
self.assertEqual(result.lshape, (4, 5))
self.assertEqual(result.split, None)
self.assertTrue(
(result._DNDarray__array == data._DNDarray__array.argmin(0)).all())
# 3D local tensor, minor axis
result = ht.argmin(data, axis=-1, keepdim=True)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.int64)
self.assertEqual(result._DNDarray__array.dtype, torch.int64)
self.assertEqual(result.shape, (3, 4, 1))
self.assertEqual(result.lshape, (3, 4, 1))
self.assertEqual(result.split, None)
self.assertTrue(
(result._DNDarray__array == data._DNDarray__array.argmin(
-1, keepdim=True)).all())
# 2D split tensor, along the axis
data = ht.array(ht.random.randn(4, 5), is_split=0)
result = ht.argmin(data, axis=1)
expected = torch.argmin(data._DNDarray__array, dim=1)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.int64)
self.assertEqual(result._DNDarray__array.dtype, torch.int64)
self.assertEqual(result.shape, (ht.MPI_WORLD.size * 4, ))
self.assertEqual(result.lshape, (4, ))
self.assertEqual(result.split, 0)
self.assertTrue((result._DNDarray__array == expected).all())
# 2D split tensor, across the axis
size = ht.MPI_WORLD.size * 2
data = ht.triu(ht.ones((size, size), split=0), k=1)
result = ht.argmin(data, axis=0)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(result.dtype, ht.int64)
self.assertEqual(result._DNDarray__array.dtype, torch.int64)
self.assertEqual(result.shape, (size, ))
self.assertEqual(result.lshape, (size, ))
self.assertEqual(result.split, None)
# skip test on gpu; argmin works different
if not (torch.cuda.is_available() and result.device == ht.gpu):
self.assertTrue((result._DNDarray__array != 0).all())
# 2D split tensor, across the axis, output tensor
size = ht.MPI_WORLD.size * 2
data = ht.triu(ht.ones((size, size), split=0), k=1)
output = ht.empty((size, ))
result = ht.argmin(data, axis=0, out=output)
self.assertIsInstance(result, ht.DNDarray)
self.assertEqual(output.dtype, ht.int64)
self.assertEqual(output._DNDarray__array.dtype, torch.int64)
self.assertEqual(output.shape, (size, ))
self.assertEqual(output.lshape, (size, ))
self.assertEqual(output.split, None)
# skip test on gpu; argmin works different
if not (torch.cuda.is_available() and output.device == ht.gpu):
self.assertTrue((output._DNDarray__array != 0).all())
# check exceptions
with self.assertRaises(TypeError):
data.argmin(axis=(0, 1))
with self.assertRaises(TypeError):
data.argmin(axis=1.1)
with self.assertRaises(TypeError):
data.argmin(axis="y")
with self.assertRaises(ValueError):
ht.argmin(data, axis=-4)