aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

implementing keepdim sum

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-17 19:08:09 -03:00
parent abcd17e679
commit 5dd3f3d202
15 changed files with 119 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

18
tests/test_operations.py
View file

@ -249,6 +249,24 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_sum_axis_keepdim(self):
"""
Test summation of a tensor along a specific axis with keepdim=True
"""
norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
norch_result = norch_tensor.sum(axis=1, keepdim=True)
torch_result = utils.to_torch(norch_result).to(self.device)
torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
torch_expected = torch.sum(torch_tensor, dim=1, keepdim=True)
print(torch_result, '\n', torch_expected)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_transpose_T(self):
"""
Test transposition of a tensor: tensor.T