Added various PyTorch wrapper nodes to generate tensors.
This commit is contained in:
HowToSD
@@ -0,0 +1,71 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_full import PtFull
|
||||
from pytorch_wrapper.utils import DTYPE_MAPPING
|
||||
|
||||
|
||||
class TestPtFull(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtFull()
|
||||
|
||||
# Tensor size to test
|
||||
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
|
||||
self.expected_shape = (2, 3)
|
||||
|
||||
# Values to test (using representative values for each type)
|
||||
self.test_values = {
|
||||
"float32": "1.5",
|
||||
"float16": "-2.3",
|
||||
"bfloat16": "3.7",
|
||||
"float64": "4.9",
|
||||
"uint8": "255", # Max value for uint8
|
||||
"int8": "-128", # Min value for int8
|
||||
"int16": "32767", # Max value for int16
|
||||
"int32": "-2147483648", # Min value for int32
|
||||
"int64": "9223372036854775807", # Max value for int64
|
||||
"bool": "True"
|
||||
}
|
||||
|
||||
def test_full_tensor_all_dtypes(self):
|
||||
"""Test tensor creation with a specified value for all supported data types."""
|
||||
for data_type, torch_dtype in DTYPE_MAPPING.items():
|
||||
value_str = self.test_values[data_type]
|
||||
|
||||
with self.subTest(data_type=data_type, value=value_str):
|
||||
result, = self.node.f(value_str, self.size_str, data_type)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has correct dtype
|
||||
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
|
||||
|
||||
# Convert value_str to correct type
|
||||
if torch_dtype in {torch.float32, torch.float16, torch.bfloat16, torch.float64}:
|
||||
expected_value = float(value_str)
|
||||
elif torch_dtype in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}:
|
||||
expected_value = int(value_str)
|
||||
elif torch_dtype == torch.bool:
|
||||
expected_value = value_str.lower() == "true"
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {data_type}")
|
||||
|
||||
# Ensure all elements match the expected value
|
||||
expected = torch.full(self.expected_shape, expected_value, dtype=torch_dtype)
|
||||
self.assertTrue(torch.equal(result, expected), f"Tensor values do not match expected tensor for dtype {data_type}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,44 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_ones import PtOnes # Importing the PtOnes class
|
||||
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
|
||||
|
||||
|
||||
class TestPtOnes(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtOnes()
|
||||
|
||||
# Tensor size to test
|
||||
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
|
||||
self.expected_shape = (2, 3)
|
||||
|
||||
def test_ones_tensor_all_dtypes(self):
|
||||
"""Test tensor creation with ones for all supported data types."""
|
||||
for data_type, torch_dtype in DTYPE_MAPPING.items():
|
||||
with self.subTest(data_type=data_type):
|
||||
result, = self.node.f(self.size_str, data_type)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has correct dtype
|
||||
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
|
||||
|
||||
# Ensure all elements are ones
|
||||
expected = torch.ones(self.expected_shape, dtype=torch_dtype)
|
||||
self.assertTrue(torch.equal(result, expected), f"Tensor values do not match expected ones tensor for dtype {data_type}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,59 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_permute import PtPermute
|
||||
|
||||
|
||||
class TestPtPermute(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtPermute()
|
||||
|
||||
# Define test cases: (input_tensor, new_axes, expected_shape)
|
||||
self.test_cases = [
|
||||
# Case 1: Basic 2D permutation
|
||||
(torch.ones(2, 3), "[1, 0]", (3, 2)), # (2,3) -> (3,2)
|
||||
|
||||
# Case 2: Permute dimensions of a 3D tensor
|
||||
(torch.ones(2, 3, 4), "[2, 0, 1]", (4, 2, 3)), # (2,3,4) -> (4,2,3)
|
||||
|
||||
# Case 3: Identity permutation (no change)
|
||||
(torch.ones(2, 3, 4), "[0, 1, 2]", (2, 3, 4)), # (2,3,4) -> (2,3,4)
|
||||
|
||||
# Case 4: Permute 4D tensor (like a channel-first transformation)
|
||||
(torch.ones(2, 3, 96, 32), "[0, 3, 1, 2]", (2, 32, 3, 96)), # (2,3,96,32) -> (2,32,3,96)
|
||||
|
||||
# Case 5: Swapping first and last dimensions
|
||||
(torch.ones(5, 10, 15), "[2, 1, 0]", (15, 10, 5)), # (5,10,15) -> (15,10,5)
|
||||
|
||||
# Case 6: Edge case with a single-dimension tensor
|
||||
(torch.ones(10,), "[0]", (10,)), # (10,) -> (10,) (No change)
|
||||
|
||||
# Case 7: Larger 5D tensor permutation
|
||||
(torch.ones(2, 3, 4, 5, 6), "[4, 3, 2, 1, 0]", (6, 5, 4, 3, 2)), # Reverse all dimensions
|
||||
]
|
||||
|
||||
def test_permute(self):
|
||||
"""Test tensor permutation operation with various dimension orders."""
|
||||
for input_tensor, new_axes, expected_shape in self.test_cases:
|
||||
with self.subTest(input_shape=input_tensor.shape, new_axes=new_axes):
|
||||
result, = self.node.f(input_tensor, new_axes)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has the same dtype
|
||||
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after permute operation")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_neg import PtNeg # Updated import to PtNeg
|
||||
|
||||
|
||||
class TestPtNeg(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtNeg()
|
||||
|
||||
# Test cases: (Tensor A, Expected result)
|
||||
self.test_cases = [
|
||||
(torch.tensor([1, -2, 3]), torch.tensor([-1, 2, -3])),
|
||||
(torch.tensor([0, 5, -7]), torch.tensor([0, -5, 7])),
|
||||
(torch.ones(3, 3), torch.full((3, 3), -1.0)), # Negation of ones
|
||||
(torch.zeros(2, 2), torch.zeros(2, 2)), # Negation of zeros
|
||||
(torch.tensor([[5], [-10]]), torch.tensor([[-5], [10]])), # Column vector case
|
||||
]
|
||||
|
||||
def test_neg(self):
|
||||
"""Test tensor element-wise negation using torch.neg()."""
|
||||
for tens_a, expected in self.test_cases:
|
||||
with self.subTest(tens_a=tens_a):
|
||||
result, = self.node.f(tens_a)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result matches expected negation
|
||||
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,46 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_rand import PtRand # Importing the PtRand class
|
||||
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
|
||||
|
||||
|
||||
class TestPtRand(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtRand()
|
||||
|
||||
# Tensor size to test
|
||||
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
|
||||
self.expected_shape = (2, 3)
|
||||
|
||||
# Only floating-point types are valid for torch.rand()
|
||||
self.valid_dtypes = {k: v for k, v in DTYPE_MAPPING.items() if v in {torch.float32, torch.float16, torch.bfloat16, torch.float64}}
|
||||
|
||||
def test_rand_tensor_all_floating_dtypes(self):
|
||||
"""Test random tensor creation for all supported floating-point data types."""
|
||||
for data_type, torch_dtype in self.valid_dtypes.items():
|
||||
with self.subTest(data_type=data_type):
|
||||
result, = self.node.f(self.size_str, data_type)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has correct dtype
|
||||
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
|
||||
|
||||
# Ensure all elements are within range [0,1)
|
||||
self.assertTrue(torch.all(result >= 0) and torch.all(result < 1), f"Tensor values are not in the range [0,1) for dtype {data_type}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,61 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
import ast
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_rand_int import PtRandInt # Importing the PtRandInt class
|
||||
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
|
||||
|
||||
|
||||
class TestPtRandInt(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtRandInt()
|
||||
|
||||
# Tensor size to test
|
||||
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
|
||||
self.expected_shape = (2, 3)
|
||||
|
||||
# Only integer types are valid for torch.randint()
|
||||
self.valid_dtypes = {
|
||||
k: v for k, v in DTYPE_MAPPING.items() if v in {torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8}
|
||||
}
|
||||
|
||||
# Define min and max values for each integer type
|
||||
self.test_ranges = {
|
||||
"int8": (-128, 127),
|
||||
"int16": (-32768, 32767),
|
||||
"int32": (-2147483648, 2147483647),
|
||||
"int64": (-9223372036854775808, 9223372036854775807),
|
||||
"uint8": (0, 255),
|
||||
}
|
||||
|
||||
def test_randint_tensor_all_integer_dtypes(self):
|
||||
"""Test random integer tensor creation for all supported integer data types."""
|
||||
for data_type, torch_dtype in self.valid_dtypes.items():
|
||||
min_value, max_value = self.test_ranges[data_type]
|
||||
|
||||
with self.subTest(data_type=data_type, min_value=min_value, max_value=max_value):
|
||||
result, = self.node.f(min_value, max_value, self.size_str, data_type)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has correct dtype
|
||||
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
|
||||
|
||||
# Ensure all elements are within range [min_value, max_value)
|
||||
self.assertTrue(torch.all(result >= min_value), f"Tensor contains values less than min_value for dtype {data_type}")
|
||||
self.assertTrue(torch.all(result < max_value), f"Tensor contains values greater than or equal to max_value for dtype {data_type}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,57 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_squeeze import PtSqueeze # Importing the PtSqueeze class
|
||||
|
||||
|
||||
class TestPtSqueeze(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtSqueeze()
|
||||
|
||||
# Define test cases: (input_tensor, dim, expected_shape)
|
||||
self.test_cases = [
|
||||
# Case 1: Squeezing a singleton dimension
|
||||
(torch.ones(1, 3, 4), 0, (3, 4)), # (1,3,4) -> (3,4)
|
||||
(torch.ones(2, 1, 4), 1, (2, 4)), # (2,1,4) -> (2,4)
|
||||
|
||||
# Case 2: No change if dim is not a singleton
|
||||
(torch.ones(2, 3, 4), 1, (2, 3, 4)), # (2,3,4) -> (2,3,4)
|
||||
|
||||
# Case 3: Multiple singleton dimensions, squeezing one
|
||||
(torch.ones(1, 2, 1, 4), 2, (1, 2, 4)), # (1,2,1,4) -> (1,2,4)
|
||||
|
||||
# Case 4: Squeezing a negative index
|
||||
(torch.ones(2, 3, 1), -1, (2, 3)), # (2,3,1) -> (2,3)
|
||||
|
||||
# Case 5: Squeezing all size 1 dimensions
|
||||
(torch.ones(1, 1, 1), 0, (1, 1)), # (1,1,1) -> (1,1)
|
||||
(torch.ones(1, 1, 1), 1, (1, 1)), # (1,1,1) -> (1,1)
|
||||
(torch.ones(1, 1, 1), 2, (1, 1)), # (1,1,1) -> (1,1)
|
||||
]
|
||||
|
||||
def test_squeeze(self):
|
||||
"""Test tensor squeeze operation on various dimensions."""
|
||||
for input_tensor, dim, expected_shape in self.test_cases:
|
||||
with self.subTest(input_shape=input_tensor.shape, dim=dim):
|
||||
result, = self.node.f(input_tensor, dim)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has the same dtype
|
||||
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after squeeze operation")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,57 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_unsqueeze import PtUnsqueeze # Importing the PtUnsqueeze class
|
||||
|
||||
|
||||
class TestPtUnsqueeze(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtUnsqueeze()
|
||||
|
||||
# Define test cases: (input_tensor, dim, expected_shape)
|
||||
self.test_cases = [
|
||||
# Case 1: Adding a singleton dimension at different positions
|
||||
(torch.ones(3, 4), 0, (1, 3, 4)), # (3,4) -> (1,3,4)
|
||||
(torch.ones(3, 4), 1, (3, 1, 4)), # (3,4) -> (3,1,4)
|
||||
(torch.ones(3, 4), 2, (3, 4, 1)), # (3,4) -> (3,4,1)
|
||||
|
||||
# Case 2: Adding a singleton dimension at negative indices
|
||||
(torch.ones(3, 4), -1, (3, 4, 1)), # (3,4) -> (3,4,1)
|
||||
(torch.ones(3, 4), -2, (3, 1, 4)), # (3,4) -> (3,1,4)
|
||||
(torch.ones(3, 4), -3, (1, 3, 4)), # (3,4) -> (1,3,4)
|
||||
|
||||
# Case 3: Adding multiple dimensions one by one
|
||||
(torch.ones(3,), 0, (1, 3)), # (3,) -> (1,3)
|
||||
(torch.ones(3,), 1, (3, 1)), # (3,) -> (3,1)
|
||||
(torch.ones(3,), -1, (3, 1)), # (3,) -> (3,1)
|
||||
|
||||
# Case 4: Expanding an already multi-dimensional tensor
|
||||
(torch.ones(2, 3, 4), 1, (2, 1, 3, 4)), # (2,3,4) -> (2,1,3,4)
|
||||
(torch.ones(2, 3, 4), -1, (2, 3, 4, 1)), # (2,3,4) -> (2,3,4,1)
|
||||
]
|
||||
|
||||
def test_unsqueeze(self):
|
||||
"""Test tensor unsqueeze operation on various dimensions."""
|
||||
for input_tensor, dim, expected_shape in self.test_cases:
|
||||
with self.subTest(input_shape=input_tensor.shape, dim=dim):
|
||||
result, = self.node.f(input_tensor, dim)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has the same dtype
|
||||
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after unsqueeze operation")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,61 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_view import PtView # Importing the PtView class
|
||||
|
||||
|
||||
class TestPtView(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtView()
|
||||
|
||||
# Define test cases: (input_tensor, new_shape, expected_shape)
|
||||
self.test_cases = [
|
||||
# Case 1: Reshape a 2D tensor
|
||||
(torch.ones(2, 3), "[3, 2]", (3, 2)), # (2,3) -> (3,2)
|
||||
|
||||
# Case 2: Reshape a 3D tensor
|
||||
(torch.ones(2, 3, 4), "[6, 4]", (6, 4)), # (2,3,4) -> (6,4)
|
||||
|
||||
# Case 3: Using -1 for automatic inference
|
||||
(torch.ones(2, 3, 4), "[2, -1]", (2, 12)), # (2,3,4) -> (2,12)
|
||||
(torch.ones(6, 4), "[-1, 4]", (6, 4)), # (6,4) -> (6,4)
|
||||
|
||||
# Case 4: Flattening a multi-dimensional tensor
|
||||
(torch.ones(2, 3, 4), "[-1]", (24,)), # (2,3,4) -> (24,)
|
||||
|
||||
# Case 5: Expanding dimensions
|
||||
(torch.ones(6,), "[2, 3]", (2, 3)), # (6,) -> (2,3)
|
||||
|
||||
# Case 6: Reshaping a 4D tensor
|
||||
(torch.ones(2, 3, 4, 5), "[6, -1, 5]", (6, 4, 5)), # (2,3,4,5) -> (6,4,5)
|
||||
|
||||
# Case 7: Edge case with 1D tensor
|
||||
(torch.ones(10,), "[10, 1]", (10, 1)), # (10,) -> (10,1)
|
||||
]
|
||||
|
||||
def test_view(self):
|
||||
"""Test tensor reshape operation with various shape orders."""
|
||||
for input_tensor, new_shape, expected_shape in self.test_cases:
|
||||
with self.subTest(input_shape=input_tensor.shape, new_shape=new_shape):
|
||||
result, = self.node.f(input_tensor, new_shape)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has the same dtype
|
||||
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after view operation")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,44 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
import ast
|
||||
|
||||
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
||||
sys.path.append(MODULE_ROOT)
|
||||
|
||||
from pytorch_wrapper.pt_zeros import PtZeros # Importing the PtZeros class
|
||||
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
|
||||
|
||||
|
||||
class TestPtZeros(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test instance."""
|
||||
self.node = PtZeros()
|
||||
|
||||
# Tensor size to test
|
||||
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
|
||||
self.expected_shape = (2, 3)
|
||||
|
||||
def test_zeros_tensor_all_dtypes(self):
|
||||
"""Test tensor creation with zeros for all supported data types."""
|
||||
for data_type, torch_dtype in DTYPE_MAPPING.items():
|
||||
with self.subTest(data_type=data_type):
|
||||
result, = self.node.f(self.size_str, data_type)
|
||||
|
||||
# Ensure return value is a tensor
|
||||
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
|
||||
|
||||
# Ensure result has correct shape
|
||||
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
|
||||
|
||||
# Ensure result has correct dtype
|
||||
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
|
||||
|
||||
# Ensure all elements are zeros
|
||||
expected = torch.zeros(self.expected_shape, dtype=torch_dtype)
|
||||
self.assertTrue(torch.equal(result, expected), f"Tensor values do not match expected zeros tensor for dtype {data_type}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user