nodes/HowToSD-ComfyUI-Data-Analysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added more PyTorch wrapper nodes.

Browse Source
This commit is contained in:
HowToSD
2025-02-27 23:52:21 -08:00
parent b0782fb422
commit b22193dc2b
58 changed files with 2753 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pytorch_wrapper/pt_argmax_test.py
+39
View File
@@ -0,0 +1,39 @@
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_argmax import PtArgmax
class TestPtArgmax(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtArgmax()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
(torch.tensor([[1, 3, 2], [4, 0, 5]]), "0", False, torch.tensor([1, 0, 1])),
(torch.tensor([[1, 3, 2], [4, 0, 5]]), "1", True, torch.tensor([[1], [2]])),
(torch.tensor([[10, 3, 2], [123, 0, 5]]), "", False, torch.tensor(3)),
]
def test_argmax(self):
"""Test torch.argmax behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected indices
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_argmin_test.py
+39
View File
@@ -0,0 +1,39 @@
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_argmin import PtArgmin
class TestPtArgmin(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtArgmin()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
(torch.tensor([[1, 3, 2], [4, 0, 5]]), "0", False, torch.tensor([0, 1, 0])),
(torch.tensor([[1, 3, 2], [4, 0, 5]]), "1", True, torch.tensor([[0], [1]])),
(torch.tensor([[1, 3, 2], [4, 0, 5]]), "", False, torch.tensor(4)),
]
def test_argmin(self):
"""Test torch.argmin behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected indices
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_gather_test.py
+48
View File
@@ -0,0 +1,48 @@
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_gather import PtGather
class TestPtGather(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtGather()
# Test cases: (Tensor, Dim, Index (as string), Expected result)
self.test_cases = [
(
torch.tensor([[10, 20, 30, 40, 50], [100, 200, 300, 400, 500]]),
1,
"[[0, 4, 3, 2, 0], [4, 3, 0, 0, 0]]",
torch.tensor([[10, 50, 40, 30, 10], [500, 400, 100, 100, 100]])
),
(
torch.tensor([[1, 2], [3, 4]]),
0,
"[[1, 0], [0, 1]]",
torch.tensor([[3, 2], [1, 4]])
),
]
def test_gather(self):
"""Test torch.gather behavior."""
for tens, dim, index, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, index=index):
result, = self.node.f(tens, dim, index)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected gathered values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_index_select_test.py
+48
View File
@@ -0,0 +1,48 @@
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_index_select import PtIndexSelect
class TestPtIndexSelect(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtIndexSelect()
# Test cases: (Tensor, Dim, Index (as string), Expected result)
self.test_cases = [
(
torch.tensor([[1, 2, 3], [40, 50, 60], [700, 800, 900]]),
0,
"[0, 2]",
torch.tensor([[1, 2, 3], [700, 800, 900]])
),
(
torch.tensor([[1, 2, 3], [40, 50, 60], [700, 800, 900]]),
1,
"[1, 2]",
torch.tensor([[2, 3], [50, 60], [800, 900]])
),
]
def test_index_select(self):
"""Test torch.index_select behavior."""
for tens, dim, index, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, index=index):
result, = self.node.f(tens, dim, index)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected selected values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_masked_select_test.py
+40
View File
@@ -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_masked_select import PtMaskedSelect
class TestPtMaskedSelect(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtMaskedSelect()
# Test cases: (Tensor, Masked Tensor, Expected result)
self.test_cases = [
(
torch.tensor([[1, 2, 3], [40, 50, 60], [700, 800, 900]]),
torch.tensor([[True, False, False], [False, True, False], [False, False, True]]),
torch.tensor([1, 50, 900])
),
]
def test_masked_select(self):
"""Test torch.masked_select behavior."""
for tens, masked_tens, expected in self.test_cases:
with self.subTest(tens=tens, masked_tens=masked_tens):
result, = self.node.f(tens, masked_tens)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected selected values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_max_test.py
+61
View File
@@ -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_max import PtMax
class TestPtMax(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtMax()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"0",
False,
torch.tensor([4, 3, 5])
),
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"1",
True,
torch.tensor([[3], [5]])
),
(
torch.tensor([[[1, 5], [2, 8]], [[3, 4], [6, 7]]]),
"(0, 1)",
False,
torch.tensor([6, 8])
),
(
# Reduce across all axes
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"", # Passing empty string to reduce across all dimensions
False,
torch.tensor(5)
)
]
def test_max(self):
"""Test torch.max behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected max values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_mean_test.py
+87
View File
@@ -0,0 +1,87 @@
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_mean import PtMean
class TestPtMean(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtMean()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
# Float tensor cases
(
torch.tensor([[1.0, 3.0, 2.0], [4.0, 0.0, 5.0]], dtype=torch.float),
"0",
False,
torch.tensor([2.5, 1.5, 3.5], dtype=torch.float)
),
(
torch.tensor([[1.0, 3.0, 2.0], [4.0, 0.0, 5.0]], dtype=torch.float),
"1",
True,
torch.tensor([[2.], [3.]], dtype=torch.float)
),
(
torch.tensor([[[1.0, 5.0], [2.0, 8.0]], [[3.0, 4.0], [6.0, 7.0]]], dtype=torch.float),
"(0, 1)",
False,
torch.tensor([3.0, 6.0], dtype=torch.float)
),
(
# Reduce across all axes
torch.tensor([[1.0, 3.0, 2.0], [4.0, 0.0, 5.0]], dtype=torch.float),
"", # Passing empty string to reduce across all dimensions
False,
torch.tensor(2.5, dtype=torch.float)
),
# Complex tensor cases
(
torch.tensor([[1 + 1j, 3 + 3j, 2 + 2j], [4 + 4j, 0 + 0j, 5 + 5j]], dtype=torch.complex64),
"0",
False,
torch.tensor([2.5 + 2.5j, 1.5 + 1.5j, 3.5 + 3.5j], dtype=torch.complex64)
),
(
torch.tensor([[1 + 1j, 3 + 3j, 2 + 2j], [4 + 4j, 0 + 0j, 5 + 5j]], dtype=torch.complex64),
"1",
True,
torch.tensor([[2. + 2j], [3. + 3j]], dtype=torch.complex64)
),
(
torch.tensor([[[1 + 1j, 5 + 5j], [2 + 2j, 8 + 8j]], [[3 + 3j, 4 + 4j], [6 + 6j, 7 + 7j]]], dtype=torch.complex64),
"(0, 1)",
False,
torch.tensor([3. + 3j, 6. + 6j], dtype=torch.complex64)
),
(
# Reduce across all axes
torch.tensor([[1 + 1j, 3 + 3j, 2 + 2j], [4 + 4j, 0 + 0j, 5 + 5j]], dtype=torch.complex64),
"", # Passing empty string to reduce across all dimensions
False,
torch.tensor(2.5 + 2.5j, dtype=torch.complex64)
)
]
def test_mean(self):
"""Test torch.mean behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected mean values
self.assertTrue(torch.allclose(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_median_test.py
+66
View File
@@ -0,0 +1,66 @@
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_median import PtMedian
class TestPtMedian(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtMedian()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
(
torch.tensor([[1, 3, 2], [2, 1, 3], [4, 0, 5]]),
"0",
False,
torch.tensor([2, 1, 3])
),
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"0",
False,
torch.tensor([1, 0, 2]) # Fix: torch.median returns the lower value when between two numbers
),
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"1",
True,
torch.tensor([[2], [4]])
),
(
torch.tensor([[1, 3, 2], [4, 0, 5], [7, 8, 6], [9, 2, 10]]),
"0",
False,
torch.tensor([4, 2, 5])
),
(
torch.tensor([[1, 3, 2, 8], [4, 0, 5, 6], [7, 8, 6, 9]]),
"1",
False,
torch.tensor([2, 4, 7])
)
]
def test_median(self):
"""Test torch.median behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected median values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_min_test.py
+61
View File
@@ -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_min import PtMin
class TestPtMin(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtMin()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"0",
False,
torch.tensor([1, 0, 2])
),
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"1",
True,
torch.tensor([[1], [0]])
),
(
torch.tensor([[[1, 5], [2, 8]], [[3, 4], [6, 7]]]),
"(0, 1)",
False,
torch.tensor([1, 4])
),
(
# Reduce across all axes
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"", # Passing empty string to reduce across all dimensions
False,
torch.tensor(0)
)
]
def test_amin(self):
"""Test torch.amin behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected min values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_prod_test.py
+59
View File
@@ -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_prod import PtProd
class TestPtProd(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtProd()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
(
torch.tensor([[1, 3, 2], [2, 1, 3]]),
"0",
False,
torch.tensor([2, 3, 6])
),
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"1",
True,
torch.tensor([[6], [0]])
),
(
torch.tensor([[1, 3, 2], [4, 5, 6]]),
"0",
False,
torch.tensor([4, 15, 12])
),
(
torch.tensor([[1, 3, 2], [4, 5, 6], [7, 8, 9]]),
"1",
False,
torch.tensor([6, 120, 504])
)
]
def test_prod(self):
"""Test torch.prod behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected product values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_reshape_test.py
+60
View File
@@ -0,0 +1,60 @@
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_reshape import PtReshape # Importing the PtReshape class
class TestPtReshape(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtReshape()
# 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_reshape(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 reshape operation")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_scatter_test.py
+51
View File
@@ -0,0 +1,51 @@
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_scatter import PtScatter
class TestPtScatter(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtScatter()
# Test cases: (Target Tensor, Dim, Index Tensor (as string), Source Tensor (as string), Expected Output)
self.test_cases = [
(
torch.tensor([
[10, 11, 12],
[21, 22, 23],
[31, 32, 33],
]),
0,
"[[0, 1, 0], [1, 1, 2]]",
"[[2, 4, 6], [4, 4, 2]]",
torch.tensor([
[2, 11, 6],
[4, 4, 23],
[31, 32, 2],
])
)
]
def test_scatter(self):
"""Test torch.scatter behavior."""
for tens, dim, index, src, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, index=index, src=src):
result, = self.node.f(tens, dim, index, src)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_sqrt_test.py
+39
View File
@@ -0,0 +1,39 @@
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_sqrt import PtSqrt # Updated import to PtSqrt
class TestPtSqrt(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtSqrt()
# Test cases: (Tensor A, Expected result)
self.test_cases = [
(torch.tensor([1.0, 4.0, 9.0], dtype=torch.float32), torch.sqrt(torch.tensor([1.0, 4.0, 9.0], dtype=torch.float32))), # Perfect squares
(torch.tensor([0.5, 2.0, 3.0], dtype=torch.float32), torch.sqrt(torch.tensor([0.5, 2.0, 3.0], dtype=torch.float32))), # Common values
(torch.tensor([[1.5, 3.2], [4.7, 5.9]], dtype=torch.float32), torch.sqrt(torch.tensor([[1.5, 3.2], [4.7, 5.9]], dtype=torch.float32))), # 2D case
(torch.full((2, 2), 2.5, dtype=torch.float32), torch.sqrt(torch.full((2, 2), 2.5, dtype=torch.float32))), # Constant matrix
]
def test_sqrt(self):
"""Test tensor element-wise square root using torch.sqrt()."""
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 sqrt values
self.assertTrue(torch.allclose(result, expected, atol=1e-6), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_std_test.py
+84
View File
@@ -0,0 +1,84 @@
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_std import PtStd
class TestPtStd(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtStd()
# Define test cases with string-based dimensions
self.test_cases = [
# Sample standard deviation (Bessel's correction) along dim=0
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"0",
1,
False,
torch.std(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=0, correction=1, keepdim=False)
),
# Sample standard deviation along dim=1, with keepdim=True
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"1",
1,
True,
torch.std(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=1, correction=1, keepdim=True)
),
# Sample standard deviation across both axes (dim=(0,1))
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"(0, 1)",
1,
False,
torch.std(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=(0, 1), correction=1, keepdim=False)
),
# Sample standard deviation across all elements (dim=None) - Entire tensor
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"",
1,
False,
torch.std(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), correction=1, keepdim=False)
),
# Population standard deviation along dim=0 (correction=0)
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"0",
0,
False,
torch.std(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=0, correction=0, keepdim=False)
),
# Population standard deviation along dim=1 with keepdim=True (correction=0)
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"1",
0,
True,
torch.std(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=1, correction=0, keepdim=True)
)
]
def test_std(self):
"""Test torch.std behavior."""
for tens, dim, correction, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, correction=correction, keepdim=keepdim):
result, = self.node.f(tens, dim, correction, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected standard deviation values
self.assertTrue(torch.allclose(result, expected, atol=1e-6), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_sum_test.py
+71
View File
@@ -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_sum import PtSum
class TestPtSum(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtSum()
# Test cases: (Tensor, Dim (as string), keepdim, Expected result)
self.test_cases = [
(
torch.tensor([[1, 3, 2], [2, 1, 3]]),
"0",
False,
torch.tensor([3, 4, 5])
),
(
torch.tensor([[1, 3, 2], [4, 0, 5]]),
"1",
True,
torch.tensor([[6], [9]])
),
(
torch.tensor([[1, 3, 2], [4, 5, 6]]),
"0",
False,
torch.tensor([5, 8, 8])
),
(
torch.tensor([[1, 3, 2], [4, 5, 6], [7, 8, 9]]),
"1",
False,
torch.tensor([6, 15, 24])
),
(
torch.tensor([[1, 2], [3, 4]]),
"",
False,
torch.tensor(10) # Sum of all elements
),
(
torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
"(0, 1)",
False,
torch.tensor([16, 20]) # Sum over dimensions (0, 1)
)
]
def test_sum(self):
"""Test torch.sum behavior."""
for tens, dim, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, keepdim=keepdim):
result, = self.node.f(tens, dim, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected sum values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_var_test.py
+83
View File
@@ -0,0 +1,83 @@
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_var import PtVar
class TestPtVar(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtVar()
# Define test cases with string-based dimensions
self.test_cases = [
# Sample variance (Bessel's correction) along dim=0
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"0",
1,
False,
torch.var(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=0, correction=1, keepdim=False)
),
# Sample variance along dim=1, with keepdim=True
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"1",
1,
True,
torch.var(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=1, correction=1, keepdim=True)
),
# Sample variance across both axes (dim=(0,1))
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"(0, 1)",
1,
False,
torch.var(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=(0, 1), correction=1, keepdim=False)
),
# Sample variance across all elements (dim=None) - Entire tensor
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"",
1,
False,
torch.var(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), correction=1, keepdim=False)
),
# Population variance along dim=0 (correction=0)
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"0",
0,
False,
torch.var(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=0, correction=0, keepdim=False)
),
# Population variance along dim=1 with keepdim=True (correction=0)
(
torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
"1",
0,
True,
torch.var(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dim=1, correction=0, keepdim=True)
)
]
def test_var(self):
"""Test torch.var behavior."""
for tens, dim, correction, keepdim, expected in self.test_cases:
with self.subTest(tens=tens, dim=dim, correction=correction, keepdim=keepdim):
result, = self.node.f(tens, dim, correction, keepdim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected variance values
self.assertTrue(torch.allclose(result, expected, atol=1e-6), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_where_test.py
+59
View File
@@ -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_where import PtWhere
class TestPtWhere(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtWhere()
# Test cases: (Condition Tensor, True Tensor, False Tensor, Expected Output)
self.test_cases = [
(
torch.tensor([[True, False], [False, True]]),
torch.tensor([[1, 2], [3, 4]]),
torch.tensor([[5, 6], [7, 8]]),
torch.tensor([[1, 6], [7, 4]])
),
(
torch.tensor([True, False, True]),
torch.tensor([10, 20, 30]),
torch.tensor([100, 200, 300]),
torch.tensor([10, 200, 30])
),
(
torch.tensor([[False, False], [True, True]]),
torch.tensor([[1.5, 2.5], [3.5, 4.5]]),
torch.tensor([[5.5, 6.5], [7.5, 8.5]]),
torch.tensor([[5.5, 6.5], [3.5, 4.5]])
),
(
torch.tensor([[[True, False], [False, True]], [[False, True], [True, False]]]),
torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
torch.tensor([[[9, 10], [11, 12]], [[13, 14], [15, 16]]]),
torch.tensor([[[1, 10], [11, 4]], [[13, 6], [7, 16]]])
),
]
def test_where(self):
"""Test torch.where behavior."""
for condition_tens, true_tens, false_tens, expected in self.test_cases:
with self.subTest(condition_tens=condition_tens, true_tens=true_tens, false_tens=false_tens):
result, = self.node.f(condition_tens, true_tens, false_tens)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected values
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_zeros_test.py
+1 -1
View File
@@ -2,7 +2,7 @@ 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")
tests/pytorch_wrapper/utils_test.py
+54
View File
@@ -0,0 +1,54 @@
import os
import sys
import unittest
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.utils import str_to_dim
class TestStrToDim(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
pass
def test_valid_integer(self):
"""Test parsing a valid integer string."""
self.assertEqual(str_to_dim("5"), 5)
def test_valid_tuple(self):
"""Test parsing a valid tuple string."""
self.assertEqual(str_to_dim("(1, 2, 3)"), (1, 2, 3))
def test_valid_list(self):
"""Test parsing a valid list string (should be converted to tuple)."""
self.assertEqual(str_to_dim("[4, 5, 6]"), (4, 5, 6))
def test_empty_string(self):
"""Test empty string input should return None."""
self.assertIsNone(str_to_dim(""))
def test_invalid_string(self):
"""Test invalid string input should raise ValueError."""
with self.assertRaises(ValueError):
str_to_dim("invalid")
def test_mixed_type_list(self):
"""Test a list with mixed types should raise ValueError."""
with self.assertRaises(ValueError):
str_to_dim("[1, 'a', 3]")
def test_nested_tuple(self):
"""Test a nested tuple should raise ValueError."""
with self.assertRaises(ValueError):
str_to_dim("((1, 2), 3)")
def test_dict_input(self):
"""Test a dictionary input should raise ValueError."""
with self.assertRaises(ValueError):
str_to_dim("{'a': 1, 'b': 2}")
if __name__ == "__main__":
unittest.main()