import unittest import norch from norch.utils import utils_unittests as utils import torch import sys import os class TestTensorOperations(unittest.TestCase): def setUp(self): self.device = os.environ.get('device') if self.device is None or self.device != 'cuda': self.device = 'cpu' print(f"Running tests on: {self.device}") def test_creation_and_conversion(self): """ Test creation and convertion of norch tensor to pytorch """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_tensor = utils.to_torch(norch_tensor) self.assertTrue(torch.is_tensor(torch_tensor)) def test_addition(self): """ Test addition two tensors: tensor1 + tensor2 """ norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) norch_result = norch_tensor1 + norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) torch_expected = torch_tensor1 + torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_addition_broadcasted(self): """ Test addition of two tensors with broadcasting: tensor1 + tensor2 """ norch_tensor1 = norch.Tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device) # Shape (1, 2, 3) norch_tensor2 = norch.Tensor([1, 1, 1]).to(self.device) # Shape (3) norch_result = norch_tensor1 + norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device) # Shape (1, 2, 3) torch_tensor2 = torch.tensor([1, 1, 1]).to(self.device) # Shape (3) torch_expected = torch_tensor1 + torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) # Shape (1, 2, 3) norch_tensor2 = norch.Tensor([[10, 10], [5, 6]]).to(self.device) # Shape (3) norch_result = norch_tensor1 + norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) # Shape (1, 2, 3) torch_tensor2 = torch.tensor([[[10, 10], [5, 6]]]).to(self.device) # Shape (3) torch_expected = torch_tensor1 + torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) # reversed order broadcasting norch_tensor1 = norch.Tensor([[0, 2]]).to(self.device) norch_tensor2 = norch.Tensor([[3, 4], [5, -1]]).to(self.device) norch_result = norch_tensor1 + norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[0, 2]]).to(self.device) torch_tensor2 = torch.tensor([[3, 4], [5, -1]]).to(self.device) torch_expected = torch_tensor1 + torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) norch_result = norch_tensor2 + norch_tensor1 torch_expected = torch_tensor2 + torch_tensor1 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_subtraction(self): """ Test subtraction of two tensors: tensor1 - tensor2 """ norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) norch_result = norch_tensor1 - norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) torch_expected = torch_tensor1 - torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_broadcasting_subtraction(self): """ Test subtraction of two tensors with broadcasting: tensor1 - tensor2 """ norch_tensor1 = norch.Tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device) # Shape (1, 2, 3) norch_tensor2 = norch.Tensor([1, 1, 1]).to(self.device) # Shape (3) norch_result = norch_tensor1 - norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device) # Shape (1, 2, 3) torch_tensor2 = torch.tensor([1, 1, 1]).to(self.device) # Shape (3) torch_expected = torch_tensor1 - torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) # reversed order broadcasting norch_result = norch_tensor2 - norch_tensor1 torch_result = utils.to_torch(norch_result).to(self.device) torch_expected = torch_tensor2 - torch_tensor1 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_division_by_scalar(self): """ Test division of a tensor by a scalar: tensor / scalar """ norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) scalar = 2 norch_result = norch_tensor / scalar torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) torch_expected = torch_tensor / scalar self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_scalar_division_by_tensor(self): """ Test scalar division by a tensor: scalar / tensor """ scalar = 10 norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) norch_result = scalar / norch_tensor torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) torch_expected = scalar / torch_tensor self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_matrix_multiplication(self): """ Test matrix multiplication: tensor1 @ tensor2 """ norch_tensor1 = norch.Tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[1., 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) norch_result = norch_tensor1 @ norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[1., 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) torch_expected = torch_tensor1 @ torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_elementwise_multiplication_by_scalar(self): """ Test elementwise multiplication of a tensor by a scalar: tensor * scalar """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) scalar = 2 norch_result = norch_tensor * scalar 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_tensor * scalar self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_elementwise_multiplication_by_tensor(self): """ Test elementwise multiplication of two tensors: tensor1 * tensor2 """ norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]]).to(self.device) norch_result = norch_tensor1 * norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]]).to(self.device) torch_expected = torch_tensor1 * torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_reshape(self): """ Test reshaping of a tensor: tensor.reshape(shape) """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) new_shape = [2, 4] norch_result = norch_tensor.reshape(new_shape) 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_tensor.reshape(new_shape) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_unsqueeze(self): """ Test unsqueeze operation on a tensor """ norch_tensor = norch.Tensor([[1, 2], [3, 4]]).to(self.device) # Unsqueeze at dim=0 norch_unsqueeze_0 = norch_tensor.unsqueeze(0) torch_unsqueeze_0 = utils.to_torch(norch_unsqueeze_0).to(self.device) torch_tensor = torch.tensor([[1, 2], [3, 4]]).to(self.device) torch_expected_0 = torch_tensor.unsqueeze(0) self.assertTrue(utils.compare_torch(torch_unsqueeze_0, torch_expected_0)) # Unsqueeze at dim=1 norch_unsqueeze_1 = norch_tensor.unsqueeze(1) torch_unsqueeze_1 = utils.to_torch(norch_unsqueeze_1).to(self.device) torch_expected_1 = torch_tensor.unsqueeze(1) self.assertTrue(utils.compare_torch(torch_unsqueeze_1, torch_expected_1)) # Unsqueeze at dim=2 norch_unsqueeze_2 = norch_tensor.unsqueeze(2) torch_unsqueeze_2 = utils.to_torch(norch_unsqueeze_2).to(self.device) torch_expected_2 = torch_tensor.unsqueeze(2) self.assertTrue(utils.compare_torch(torch_unsqueeze_2, torch_expected_2)) # Unsqueeze at dim=-1 norch_unsqueeze_neg_1 = norch_tensor.unsqueeze(-1) torch_unsqueeze_neg_1 = utils.to_torch(norch_unsqueeze_neg_1).to(self.device) torch_expected_neg_1 = torch_tensor.unsqueeze(-1) self.assertTrue(utils.compare_torch(torch_unsqueeze_neg_1, torch_expected_neg_1)) # Unsqueeze at dim=-2 norch_unsqueeze_neg_2 = norch_tensor.unsqueeze(-2) torch_unsqueeze_neg_2 = utils.to_torch(norch_unsqueeze_neg_2).to(self.device) torch_expected_neg_2 = torch_tensor.unsqueeze(-2) self.assertTrue(utils.compare_torch(torch_unsqueeze_neg_2, torch_expected_neg_2)) def test_squeeze(self): """ Test squeeze operation on a tensor """ # Create a tensor with some dimensions of size 1 norch_tensor = norch.Tensor([[[1, 2], [3, 4]]]).to(self.device) # shape [1, 2, 2] # Squeeze at dim=0 norch_squeeze_0 = norch_tensor.squeeze(0) torch_squeeze_0 = utils.to_torch(norch_squeeze_0).to(self.device) torch_tensor = torch.tensor([[[1, 2], [3, 4]]]).to(self.device) torch_expected_0 = torch_tensor.squeeze(0) self.assertTrue(utils.compare_torch(torch_squeeze_0, torch_expected_0)) # Create a tensor with a dimension of size 1 in the middle norch_tensor_middle_1 = norch.Tensor([[[1, 2]], [[3, 4]]]).to(self.device) # shape [2, 1, 2] # Squeeze at dim=1 norch_squeeze_1 = norch_tensor_middle_1.squeeze(1) torch_squeeze_1 = utils.to_torch(norch_squeeze_1).to(self.device) torch_tensor_middle_1 = torch.tensor([[[1, 2]], [[3, 4]]]).to(self.device) torch_expected_1 = torch_tensor_middle_1.squeeze(1) self.assertTrue(utils.compare_torch(torch_squeeze_1, torch_expected_1)) # Squeeze at dim=-2 (same as dim=1 in this case) norch_squeeze_neg_2 = norch_tensor_middle_1.squeeze(-2) torch_squeeze_neg_2 = utils.to_torch(norch_squeeze_neg_2).to(self.device) torch_expected_neg_2 = torch_tensor_middle_1.squeeze(-2) self.assertTrue(utils.compare_torch(torch_squeeze_neg_2, torch_expected_neg_2)) # Squeeze all dimensions of size 1 (None) norch_tensor_all_1 = norch.Tensor([[[[1, 2], [3, 4]]]]).to(self.device) # shape [1, 1, 2, 2] norch_squeeze_all = norch_tensor_all_1.squeeze() torch_squeeze_all = utils.to_torch(norch_squeeze_all).to(self.device) torch_tensor_all_1 = torch.tensor([[[[1, 2], [3, 4]]]]).to(self.device) torch_expected_all = torch_tensor_all_1.squeeze() self.assertTrue(utils.compare_torch(torch_squeeze_all, torch_expected_all)) # Squeeze no dimensions (no dimensions of size 1) norch_tensor_no_1 = norch.Tensor([[1, 2], [3, 4]]).to(self.device) # shape [2, 2] norch_squeeze_none = norch_tensor_no_1.squeeze() torch_squeeze_none = utils.to_torch(norch_squeeze_none).to(self.device) torch_tensor_no_1 = torch.tensor([[1, 2], [3, 4]]).to(self.device) torch_expected_none = torch_tensor_no_1.squeeze() self.assertTrue(utils.compare_torch(torch_squeeze_none, torch_expected_none)) def test_transpose(self): """ Test transposition of a tensor: tensor.transpose(dim1, dim2) """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) dim1, dim2 = 0, 2 norch_result = norch_tensor.transpose(dim1, dim2) 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_tensor.transpose(dim1, dim2) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_logarithm(self): """ Test elementwise logarithm of a tensor: tensor.log() """ norch_tensor = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = norch_tensor.log() 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.log(torch_tensor) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_sum(self): """ Test summation of a tensor: tensor.sum() """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = norch_tensor.sum() 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) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_sum_axis(self): """ Test summation of a tensor along a specific axis without keeping the dimensions """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = norch_tensor.sum(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.sum(torch_tensor, dim=1) 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): """ 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) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_max(self): """ Test max of a tensor: tensor.max() """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = norch_tensor.max() 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) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_max_axis(self): """ Test max of a tensor along a specific axis without keeping the dimensions """ 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)) # 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): """ Test max 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.max(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.max(torch_tensor, dim=1, keepdim=True) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_min(self): """ Test min of a tensor: tensor.min() """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = norch_tensor.min() 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) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_min_axis(self): """ Test min of a tensor along a specific axis without keeping the dimensions """ 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)) # 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): """ Test min 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.min(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.min(torch_tensor, dim=1, keepdim=True) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_1D_T(self): """ Test transposition of a 1D tensor. """ norch_tensor = norch.Tensor([1, 2, 3, 4]).to(self.device) norch_result = norch_tensor.T torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([1, 2, 3, 4]).to(self.device) torch_expected = torch_tensor.T self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_2D_T(self): """ Test transposition of a 2D tensor. """ norch_tensor = norch.Tensor([[1, 2, 3], [4, 5, 6]]).to(self.device) norch_result = norch_tensor.T torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]).to(self.device) torch_expected = torch_tensor.T self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_3D_T(self): """ Test transposition of a tensor: tensor.T """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = norch_tensor.T 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_tensor.T self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_matmul(self): """ Test 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) """ norch_tensor = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) new_shape = [2, 4] norch_reshaped = norch_tensor.reshape(new_shape) norch_result = norch_reshaped @ norch_tensor 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_tensor.reshape(new_shape) @ torch_tensor 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): """ Test transposing a tensor followed by matrix multiplication: (tensor.transpose(dim1, dim2) @ other_tensor) """ norch_tensor = norch.Tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) dim1, dim2 = 0, 2 norch_result = norch_tensor.transpose(dim1, dim2) @ norch_tensor 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_tensor.transpose(dim1, dim2) @ torch_tensor self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_add_div_matmul_then_reshape(self): """ Test a combination of operations: (tensor.sum() + other_tensor) / scalar @ another_tensor followed by reshape """ norch_tensor1 = norch.Tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[1., 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) scalar = 2 new_shape = [2, 4] norch_result = ((norch_tensor1 + norch_tensor2) / scalar) @ norch_tensor1 norch_result = norch_result.reshape(new_shape) torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[1., 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) torch_expected = ((torch_tensor1 + torch_tensor2) / scalar) @ torch_tensor1 torch_expected = torch_expected.reshape(new_shape) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_scalar_power_tensor(self): """ Test scalar power of a tensor: scalar ** tensor """ scalar = 3 norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = scalar ** norch_tensor torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_expected = scalar ** torch_tensor self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_tensor_power_scalar(self): """ Test tensor power of a scalar: tensor ** scalar """ scalar = 3 norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_result = norch_tensor ** scalar torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_expected = torch_tensor ** scalar self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_tensor_sin(self): """ Test sine function on tensor """ norch_tensor = norch.Tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) norch_result = norch_tensor.sin() torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) torch_expected = torch.sin(torch_tensor) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_tensor_cos(self): """ Test cosine function on tensor """ norch_tensor = norch.Tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) norch_result = norch_tensor.cos() torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor = torch.tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) torch_expected = torch.cos(torch_tensor) self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_equal(self): """ Test equal two tensors: tensor1.equal(tensor2) """ norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 1], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) norch_result = norch_tensor1.equal(norch_tensor2) torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 1], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) torch_expected = (torch_tensor1 == torch_tensor2).float() self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_broadcasted_equal(self): """ Test broadcasted equal two tensors: tensor1.equal(tensor2) """ norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[10, 10]], [[5, 6]]]).to(self.device) norch_result = norch_tensor1.equal(norch_tensor2) torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[10, 10]], [[5, 6]]]).to(self.device) torch_expected = (torch_tensor1 == torch_tensor2).float() self.assertTrue(utils.compare_torch(torch_result, torch_expected)) norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) norch_tensor2 = norch.Tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device) norch_result = norch_tensor1.equal(norch_tensor2) torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) torch_tensor2 = torch.tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device) torch_expected = (torch_tensor1 == torch_tensor2).float() self.assertTrue(utils.compare_torch(torch_result, torch_expected)) def test_zeros_like(self): """ Test creating a tensor of zeros with the same shape as another tensor. """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_zeros = norch_tensor.zeros_like() torch_zeros_result = utils.to_torch(norch_zeros).to(self.device) torch_tensor_expected = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_zeros_expected = torch.zeros_like(torch_tensor_expected) self.assertTrue(utils.compare_torch(torch_zeros_result, torch_zeros_expected)) def test_ones_like(self): """ Test creating a tensor of ones with the same shape as another tensor. """ norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) norch_ones = norch_tensor.ones_like() torch_ones_result = utils.to_torch(norch_ones).to(self.device) torch_tensor_expected = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) torch_ones_expected = torch.ones_like(torch_tensor_expected) self.assertTrue(utils.compare_torch(torch_ones_result, torch_ones_expected)) if __name__ == '__main__': unittest.main()