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

add new test matmul cases broadcasted and batched

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-16 20:34:06 -03:00
parent b3c0543114
commit 156ec9e3e3
3 changed files with 131 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
norch/__pycache__/tensor.cpython-38.pyc
View file

Binary file not shown.

66
tests/test_autograd.py
View file

@ -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

66
tests/test_operations.py
View file

@ -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):
"""