def test_nonzero(self):
num_src = 12
types = [
'torch.ByteTensor',
'torch.CharTensor',
'torch.ShortTensor',
'torch.IntTensor',
'torch.FloatTensor',
'torch.DoubleTensor',
'torch.LongTensor',
]
shapes = [
torch.Size((12,)),
torch.Size((12, 1)),
torch.Size((1, 12)),
torch.Size((6, 2)),
torch.Size((3, 2, 2)),
]
for t in types:
while True:
tensor = torch.rand(num_src).mul(2).floor().type(t)
if tensor.sum() > 0:
break
for shape in shapes:
tensor = tensor.clone().resize_(shape)
dst1 = torch.nonzero(tensor)
dst2 = tensor.nonzero()
dst3 = torch.LongTensor()
torch.nonzero(tensor, out=dst3)
if len(shape) == 1:
dst = []
for i in range(num_src):
if tensor[i] != 0:
dst += [i]
self.assertEqual(dst1.select(1, 0), torch.LongTensor(dst), 0)
self.assertEqual(dst2.select(1, 0), torch.LongTensor(dst), 0)
self.assertEqual(dst3.select(1, 0), torch.LongTensor(dst), 0)
elif len(shape) == 2:
# This test will allow through some False positives. It only checks
# that the elements flagged positive are indeed non-zero.
for i in range(dst1.size(0)):
self.assertNotEqual(tensor[dst1[i, 0], dst1[i, 1]], 0)
elif len(shape) == 3:
# This test will allow through some False positives. It only checks
# that the elements flagged positive are indeed non-zero.
for i in range(dst1.size(0)):
self.assertNotEqual(tensor[dst1[i, 0], dst1[i, 1], dst1[i, 2]], 0)
评论列表
文章目录
