softmax operation

This commit is contained in:
lucasdelimanogueira 2024-05-21 01:22:38 -03:00
parent fd831f950d
commit a39bd59a3f
9 changed files with 67 additions and 21 deletions

View file

@ -95,13 +95,13 @@ class TestNNModuleActivationFn(unittest.TestCase):
"""
# Test different axes
axes = [None, 0, 1, -1]
axes = [0, 1, 2, -1]
# Define the input tensors for different test cases
test_cases = [
(norch.Tensor([[1., 2, 3], [4, 5, 6]]), torch.tensor([[1., 2, 3], [4, 5, 6]])),
(norch.Tensor([[1., -1, 0], [2, -2, 0]]), torch.tensor([[1., -1, 0], [2, -2, 0]])),
(norch.Tensor([[0., 0, 0], [0, 0, 0]]), torch.tensor([[0., 0, 0], [0, 0, 0]]))
(norch.Tensor([[[1., 2, 3], [4, 5, 6]]]), torch.tensor([[[1., 2, 3], [4, 5, 6]]])),
(norch.Tensor([[[1., -1, 0], [2, -2, 0]]]), torch.tensor([[[1., -1, 0], [2, -2, 0]]])),
(norch.Tensor([[[0., 0, 0], [0, 0, 0]]]), torch.tensor([[[0., 0, 0], [0, 0, 0]]]))
]
for dim in axes:
@ -121,7 +121,5 @@ class TestNNModuleActivationFn(unittest.TestCase):
softmax_torch_expected = softmax_fn_torch.forward(torch_input)
# Compare the results
print(softmax_torch_result)
print(softmax_torch_expected)
self.assertTrue(utils.compare_torch(softmax_torch_result, softmax_torch_expected))

View file

@ -260,6 +260,17 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
# negative axis
norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
norch_result = norch_tensor.sum(axis=-2)
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=-2)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_sum_axis_keepdim(self):
"""
@ -301,6 +312,17 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
# negative axis
norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
norch_result = norch_tensor.max(axis=-1)
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.max(torch_tensor, dim=-1)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_max_axis_keepdim(self):
"""
@ -341,6 +363,17 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
# negative axis
norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
norch_result = norch_tensor.min(axis=-1)
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.min(torch_tensor, dim=-1)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_min_axis_keepdim(self):
"""