From 156ec9e3e30b1e993fbf74edb53bf055208555bc Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Thu, 16 May 2024 20:34:06 -0300 Subject: [PATCH] add new test matmul cases broadcasted and batched --- norch/__pycache__/tensor.cpython-38.pyc | Bin 15407 -> 15407 bytes tests/test_autograd.py | 66 +++++++++++++++++++++++- tests/test_operations.py | 66 ++++++++++++++++++++++++ 3 files changed, 131 insertions(+), 1 deletion(-) diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 5a33b4a3b38d8766ef05b6c52cb3729e188a2a2f..eeeaf809dd0adb885c280d76e140a26090808205 100644 GIT binary patch delta 19 ZcmZ2qvA%*Ul$V!_0SKyFHgc)j001~_1nmF- delta 19 ZcmZ2qvA%*Ul$V!_0SGdZH*%@k001{t1i1hJ diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 42431a8..b22eba6 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -211,7 +211,71 @@ class TestTensorAutograd(unittest.TestCase): self.assertTrue(utils.compare_torch(norch_tensor1_grad_matmul, torch_tensor1_grad_matmul)) self.assertTrue(utils.compare_torch(norch_tensor2_grad_matmul, torch_tensor2_grad_matmul)) - + def test_batched_matmul(self): + """ + Test autograd from batched matrix multiplication: BxMxP = BxNxM @ BxMxP + """ + B = 3 # Batch size + + norch_tensor1_matmul = norch.Tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)], requires_grad=True).to(self.device) + norch_tensor2_matmul = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device) + + norch_result_matmul = norch_tensor1_matmul @ norch_tensor2_matmul + norch_result_matmul_sum = norch_result_matmul.sum() # Sum over all elements + norch_result_matmul_sum.backward() + + # Convert gradients to torch tensors + norch_tensor1_grad_matmul = utils.to_torch(norch_tensor1_matmul.grad).to(self.device) + norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device) + + # Repeat the same process with torch tensors + torch_tensor1_matmul = torch.tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)], requires_grad=True).to(self.device) + torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device) + torch_result_matmul = torch.matmul(torch_tensor1_matmul, torch_tensor2_matmul) + torch_result_matmul_sum = torch_result_matmul.sum() + torch_result_matmul_sum.backward() + + # Extract gradients from torch tensors + torch_tensor1_grad_matmul = torch_tensor1_matmul.grad + torch_tensor2_grad_matmul = torch_tensor2_matmul.grad + + # Assertions to compare the gradients + self.assertTrue(utils.compare_torch(norch_tensor1_grad_matmul, torch_tensor1_grad_matmul)) + self.assertTrue(utils.compare_torch(norch_tensor2_grad_matmul, torch_tensor2_grad_matmul)) + + def test_broadcasted_batched_matmul(self): + """ + Test autograd from batched matrix multiplication: BxMxP = NxM @ BxMxP + """ + B = 3 # Batch size + + norch_tensor1_matmul = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device) + norch_tensor2_matmul = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device) + + norch_result_matmul = norch_tensor1_matmul @ norch_tensor2_matmul + norch_result_matmul_sum = norch_result_matmul.sum() # Sum over all elements + norch_result_matmul_sum.backward() + + # Convert gradients to torch tensors + norch_tensor1_grad_matmul = utils.to_torch(norch_tensor1_matmul.grad).to(self.device) + norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device) + + # Repeat the same process with torch tensors + torch_tensor1_matmul = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device) + torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device) + torch_result_matmul = torch.matmul(torch_tensor1_matmul, torch_tensor2_matmul) + torch_result_matmul_sum = torch_result_matmul.sum() + torch_result_matmul_sum.backward() + + # Extract gradients from torch tensors + torch_tensor1_grad_matmul = torch_tensor1_matmul.grad + torch_tensor2_grad_matmul = torch_tensor2_matmul.grad + + # Assertions to compare the gradients + self.assertTrue(utils.compare_torch(norch_tensor1_grad_matmul, torch_tensor1_grad_matmul)) + self.assertTrue(utils.compare_torch(norch_tensor2_grad_matmul, torch_tensor2_grad_matmul)) + + def test_elementwise_mul_scalar(self): """ Test autograd from elementwise multiplication with scalar: scalar * tensor diff --git a/tests/test_operations.py b/tests/test_operations.py index 449c531..3a6f414 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -242,6 +242,26 @@ class TestTensorOperations(unittest.TestCase): self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_matmul(self): + """ + Test batched matrix multiplication: MxP = NxM @ MxP + """ + # Creating batched tensors for Norch + norch_tensor1 = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + norch_tensor2 = norch.Tensor([[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) + + norch_result = norch_tensor1 @ norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + # Converting to PyTorch tensors for comparison + torch_tensor1 = torch.tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_tensor2 = torch.tensor([[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) + + torch_expected = torch.matmul(torch_tensor1, torch_tensor2) + + # Comparing results + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_reshape_then_matmul(self): """ Test reshaping a tensor followed by matrix multiplication: (tensor.reshape(shape) @ other_tensor) @@ -258,6 +278,52 @@ class TestTensorOperations(unittest.TestCase): self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_batched_matmul(self): + """ + Test batched matrix multiplication: BxMxP = BxNxM @ BxMxP + """ + B = 3 # Batch size + + # Creating batched tensors for Norch + norch_tensor1 = norch.Tensor([[[1, 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) + norch_tensor2 = norch.Tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + norch_result = norch_tensor1 @ norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + # Converting to PyTorch tensors for comparison + torch_tensor1 = torch.tensor([[[1, 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) + torch_tensor2 = torch.tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + torch_expected = torch.matmul(torch_tensor1, torch_tensor2) + + # Comparing results + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + def test_broadcasted_batched_matmul(self): + """ + Test broadcasted batched matrix multiplication: BxMxP = NxM @ BxMxP + """ + B = 3 # Batch size + + # Creating batched tensors for Norch + norch_tensor1 = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + norch_tensor2 = norch.Tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + norch_result = norch_tensor1 @ norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + # Converting to PyTorch tensors for comparison + torch_tensor1 = torch.tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_tensor2 = torch.tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + torch_expected = torch.matmul(torch_tensor1, torch_tensor2) + + # Comparing results + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_transpose_then_matmul(self): """