def test_empty_like(self):
# scalar
like_int = ht.empty_like(3)
self.assertIsInstance(like_int, ht.tensor)
self.assertEqual(like_int.shape, (1, ))
self.assertEqual(like_int.lshape, (1, ))
self.assertEqual(like_int.split, None)
self.assertEqual(like_int.dtype, ht.int32)
# sequence
like_str = ht.empty_like('abc')
self.assertIsInstance(like_str, ht.tensor)
self.assertEqual(like_str.shape, (3, ))
self.assertEqual(like_str.lshape, (3, ))
self.assertEqual(like_str.split, None)
self.assertEqual(like_str.dtype, ht.float32)
# elaborate tensor
ones = ht.ones((
2,
3,
), dtype=ht.uint8)
like_ones = ht.empty_like(ones)
self.assertIsInstance(like_ones, ht.tensor)
self.assertEqual(like_ones.shape, (
2,
3,
))
self.assertEqual(like_ones.lshape, (
2,
3,
))
self.assertEqual(like_ones.split, None)
self.assertEqual(like_ones.dtype, ht.uint8)
# elaborate tensor with split
ones_split = ht.ones((
2,
3,
), dtype=ht.uint8, split=0)
like_ones_split = ht.empty_like(ones_split)
self.assertIsInstance(like_ones_split, ht.tensor)
self.assertEqual(like_ones_split.shape, (
2,
3,
))
self.assertLessEqual(like_ones_split.lshape[0], 2)
self.assertEqual(like_ones_split.lshape[1], 3)
self.assertEqual(like_ones_split.split, 0)
self.assertEqual(like_ones_split.dtype, ht.uint8)
# exceptions
with self.assertRaises(TypeError):
ht.empty_like(ones, dtype='abc')
with self.assertRaises(TypeError):
ht.empty_like(ones, split='axis')
def test_sign(self):
# floats 1d
a = ht.array([-1, -0.5, 0, 0.5, 1])
signed = ht.sign(a)
comparison = ht.array([-1.0, -1, 0, 1, 1])
self.assertEqual(signed.dtype, comparison.dtype)
self.assertEqual(signed.shape, comparison.shape)
self.assertEqual(signed.device, a.device)
self.assertEqual(signed.split, a.split)
self.assertTrue(ht.equal(signed, comparison))
# complex + 2d + split
a = ht.array([[1 - 2j, -0.5 + 1j], [0, 4 + 6j]], split=0)
signed = ht.sign(a)
comparison = ht.array([[1 + 0j, -1 + 0j], [0 + 0j, 1 + 0j]], split=0)
self.assertEqual(signed.dtype, comparison.dtype)
self.assertEqual(signed.shape, comparison.shape)
self.assertEqual(signed.device, a.device)
self.assertEqual(signed.split, a.split)
self.assertTrue(ht.allclose(signed.real, comparison.real))
self.assertTrue(ht.allclose(signed.imag, comparison.imag, atol=2e-5))
# complex + split + out
a = ht.array([[1 - 2j, -0.5 + 1j], [0, 4 + 6j]], split=1)
b = ht.empty_like(a)
signed = ht.sign(a, b)
comparison = ht.array([[1 + 0j, -1 + 0j], [0 + 0j, 1 + 0j]], split=1)
self.assertIs(b, signed)
self.assertEqual(signed.dtype, comparison.dtype)
self.assertEqual(signed.shape, comparison.shape)
self.assertEqual(signed.device, a.device)
self.assertEqual(signed.split, a.split)
self.assertTrue(ht.allclose(signed.real, comparison.real))
self.assertTrue(ht.allclose(signed.imag, comparison.imag, atol=2e-5))
# zeros + 3d + complex + split
a = ht.zeros((4, 4, 4), dtype=ht.complex128, split=2)
signed = ht.sign(a)
comparison = ht.zeros((4, 4, 4), dtype=ht.complex128, split=2)
self.assertEqual(signed.dtype, comparison.dtype)
self.assertEqual(signed.shape, comparison.shape)
self.assertEqual(signed.device, a.device)
self.assertEqual(signed.split, a.split)
self.assertTrue(ht.allclose(signed.real, comparison.real))
self.assertTrue(ht.allclose(signed.imag, comparison.imag, atol=2e-5))
def test_empty_like(self):
# scalar
like_int = ht.empty_like(3, device=ht_device)
self.assertIsInstance(like_int, ht.DNDarray)
self.assertEqual(like_int.shape, (1,))
self.assertEqual(like_int.lshape, (1,))
self.assertEqual(like_int.split, None)
self.assertEqual(like_int.dtype, ht.int32)
# sequence
like_str = ht.empty_like("abc", device=ht_device)
self.assertIsInstance(like_str, ht.DNDarray)
self.assertEqual(like_str.shape, (3,))
self.assertEqual(like_str.lshape, (3,))
self.assertEqual(like_str.split, None)
self.assertEqual(like_str.dtype, ht.float32)
# elaborate tensor
ones = ht.ones((2, 3), dtype=ht.uint8, device=ht_device)
like_ones = ht.empty_like(ones, device=ht_device)
self.assertIsInstance(like_ones, ht.DNDarray)
self.assertEqual(like_ones.shape, (2, 3))
self.assertEqual(like_ones.lshape, (2, 3))
self.assertEqual(like_ones.split, None)
self.assertEqual(like_ones.dtype, ht.uint8)
# elaborate tensor with split
ones_split = ht.ones((2, 3), dtype=ht.uint8, split=0, device=ht_device)
like_ones_split = ht.empty_like(ones_split, device=ht_device)
self.assertIsInstance(like_ones_split, ht.DNDarray)
self.assertEqual(like_ones_split.shape, (2, 3))
self.assertLessEqual(like_ones_split.lshape[0], 2)
self.assertEqual(like_ones_split.lshape[1], 3)
self.assertEqual(like_ones_split.split, 0)
self.assertEqual(like_ones_split.dtype, ht.uint8)
# exceptions
with self.assertRaises(TypeError):
ht.empty_like(ones, dtype="abc", device=ht_device)
with self.assertRaises(TypeError):
ht.empty_like(ones, split="axis", device=ht_device)
def test_sort(self):
size = ht.MPI_WORLD.size
rank = ht.MPI_WORLD.rank
tensor = torch.arange(size, device=device).repeat(size).reshape(size, size)
data = ht.array(tensor, split=None, device=ht_device)
result, result_indices = ht.sort(data, axis=0, descending=True)
expected, exp_indices = torch.sort(tensor, dim=0, descending=True)
self.assertTrue(torch.equal(result._DNDarray__array, expected))
self.assertTrue(torch.equal(result_indices._DNDarray__array, exp_indices))
result, result_indices = ht.sort(data, axis=1, descending=True)
expected, exp_indices = torch.sort(tensor, dim=1, descending=True)
self.assertTrue(torch.equal(result._DNDarray__array, expected))
self.assertTrue(torch.equal(result_indices._DNDarray__array, exp_indices))
data = ht.array(tensor, split=0, device=ht_device)
exp_axis_zero = torch.arange(size, device=device).reshape(1, size)
exp_indices = torch.tensor([[rank] * size], device=device)
result, result_indices = ht.sort(data, descending=True, axis=0)
self.assertTrue(torch.equal(result._DNDarray__array, exp_axis_zero))
self.assertTrue(torch.equal(result_indices._DNDarray__array, exp_indices))
exp_axis_one, exp_indices = (
torch.arange(size, device=device).reshape(1, size).sort(dim=1, descending=True)
)
result, result_indices = ht.sort(data, descending=True, axis=1)
self.assertTrue(torch.equal(result._DNDarray__array, exp_axis_one))
self.assertTrue(torch.equal(result_indices._DNDarray__array, exp_indices))
result1 = ht.sort(data, axis=1, descending=True)
result2 = ht.sort(data, descending=True)
self.assertTrue(ht.equal(result1[0], result2[0]))
self.assertTrue(ht.equal(result1[1], result2[1]))
data = ht.array(tensor, split=1, device=ht_device)
exp_axis_zero = torch.tensor(rank, device=device).repeat(size).reshape(size, 1)
indices_axis_zero = torch.arange(size, dtype=torch.int64, device=device).reshape(size, 1)
result, result_indices = ht.sort(data, axis=0, descending=True)
self.assertTrue(torch.equal(result._DNDarray__array, exp_axis_zero))
# comparison value is only true on CPU
if result_indices._DNDarray__array.is_cuda is False:
self.assertTrue(torch.equal(result_indices._DNDarray__array, indices_axis_zero))
exp_axis_one = torch.tensor(size - rank - 1, device=device).repeat(size).reshape(size, 1)
result, result_indices = ht.sort(data, descending=True, axis=1)
self.assertTrue(torch.equal(result._DNDarray__array, exp_axis_one))
self.assertTrue(torch.equal(result_indices._DNDarray__array, exp_axis_one))
tensor = torch.tensor(
[
[[2, 8, 5], [7, 2, 3]],
[[6, 5, 2], [1, 8, 7]],
[[9, 3, 0], [1, 2, 4]],
[[8, 4, 7], [0, 8, 9]],
],
dtype=torch.int32,
device=device,
)
data = ht.array(tensor, split=0, device=ht_device)
exp_axis_zero = torch.tensor([[2, 3, 0], [0, 2, 3]], dtype=torch.int32, device=device)
if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
indices_axis_zero = torch.tensor(
[[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=device
)
else:
indices_axis_zero = torch.tensor(
[[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=device
)
result, result_indices = ht.sort(data, axis=0)
first = result[0]._DNDarray__array
first_indices = result_indices[0]._DNDarray__array
if rank == 0:
self.assertTrue(torch.equal(first, exp_axis_zero))
self.assertTrue(torch.equal(first_indices, indices_axis_zero))
data = ht.array(tensor, split=1, device=ht_device)
exp_axis_one = torch.tensor([[2, 2, 3]], dtype=torch.int32, device=device)
indices_axis_one = torch.tensor([[0, 1, 1]], dtype=torch.int32, device=device)
result, result_indices = ht.sort(data, axis=1)
first = result[0]._DNDarray__array[:1]
first_indices = result_indices[0]._DNDarray__array[:1]
if rank == 0:
self.assertTrue(torch.equal(first, exp_axis_one))
self.assertTrue(torch.equal(first_indices, indices_axis_one))
data = ht.array(tensor, split=2, device=ht_device)
exp_axis_two = torch.tensor([[2], [2]], dtype=torch.int32, device=device)
indices_axis_two = torch.tensor([[0], [1]], dtype=torch.int32, device=device)
result, result_indices = ht.sort(data, axis=2)
first = result[0]._DNDarray__array[:, :1]
first_indices = result_indices[0]._DNDarray__array[:, :1]
if rank == 0:
self.assertTrue(torch.equal(first, exp_axis_two))
self.assertTrue(torch.equal(first_indices, indices_axis_two))
#
out = ht.empty_like(data, device=ht_device)
indices = ht.sort(data, axis=2, out=out)
self.assertTrue(ht.equal(out, result))
self.assertTrue(ht.equal(indices, result_indices))
with self.assertRaises(ValueError):
ht.sort(data, axis=3)
with self.assertRaises(TypeError):
ht.sort(data, axis="1")
rank = ht.MPI_WORLD.rank
data = ht.random.randn(100, 1, split=0, device=ht_device)
result, _ = ht.sort(data, axis=0)
counts, _, _ = ht.get_comm().counts_displs_shape(data.gshape, axis=0)
for i, c in enumerate(counts):
for idx in range(c - 1):
if rank == i:
self.assertTrue(
torch.lt(
result._DNDarray__array[idx], result._DNDarray__array[idx + 1]
).all()
)