diff --git a/__init__.py b/__init__.py
index 6388e25..463a93f 100644
--- a/__init__.py
+++ b/__init__.py
@@ -17,6 +17,7 @@ from .modules.pytorch_wrapper.pt_bitwise_right_shift import PtBitwiseRightShift
from .modules.pytorch_wrapper.pt_bitwise_xor import PtBitwiseXor
from .modules.pytorch_wrapper.pt_bmm import PtBmm
from .modules.pytorch_wrapper.pt_bool_create import PtBoolCreate
+from .modules.pytorch_wrapper.pt_concat import PtConcat
from .modules.pytorch_wrapper.pt_cos import PtCos
from .modules.pytorch_wrapper.pt_cosh import PtCosh
from .modules.pytorch_wrapper.pt_crop import PtCrop
@@ -85,6 +86,7 @@ from .modules.pytorch_wrapper.pt_size_to_numpy import PtSizeToNumpy
from .modules.pytorch_wrapper.pt_size_to_string import PtSizeToString
from .modules.pytorch_wrapper.pt_sqrt import PtSqrt
from .modules.pytorch_wrapper.pt_squeeze import PtSqueeze
+from .modules.pytorch_wrapper.pt_stack import PtStack
from .modules.pytorch_wrapper.pt_std import PtStd
from .modules.pytorch_wrapper.pt_sub import PtSub
from .modules.pytorch_wrapper.pt_sum import PtSum
@@ -191,6 +193,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtBitwiseXor": PtBitwiseXor,
"PtBmm": PtBmm,
"PtBoolCreate": PtBoolCreate,
+ "PtConcat": PtConcat,
"PtCos": PtCos,
"PtCosh": PtCosh,
"PtCrop": PtCrop,
@@ -259,6 +262,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtSizeToString": PtSizeToString,
"PtSqrt": PtSqrt,
"PtSqueeze": PtSqueeze,
+ "PtStack": PtStack,
"PtStd": PtStd,
"PtSub": PtSub,
"PtSum": PtSum,
@@ -365,6 +369,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtBitwiseXor": "Pt Bitwise Xor",
"PtBmm": "Pt Bmm",
"PtBoolCreate": "Pt Bool Create",
+ "PtConcat": "Pt Concat",
"PtCos": "Pt Cos",
"PtCosh": "Pt Cosh",
"PtCrop": "Pt Crop",
@@ -433,6 +438,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtSizeToString": "Pt Size To String",
"PtSqrt": "Pt Sqrt",
"PtSqueeze": "Pt Squeeze",
+ "PtStack": "Pt Stack",
"PtStd": "Pt Std",
"PtSub": "Pt Sub",
"PtSum": "Pt Sum",
diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md
index 41e480d..f0e7537 100644
--- a/docs/reference/node_reference.md
+++ b/docs/reference/node_reference.md
@@ -221,9 +221,11 @@
## PyTorch wrapper - Transform
| Node | Description |
| --- | --- |
+| [Pt Concat](pt_concat.md) | Concatenates two PyTorch tensors. |
| [Pt Flatten](pt_flatten.md) | Flattens a PyTorch tensor into a 1D tensor. |
| [Pt Permute](pt_permute.md) | Permutes the dimensions of a PyTorch tensor according to the specified order. |
| [Pt Reshape](pt_reshape.md) | Reshapes a PyTorch tensor into a specified shape using `torch.reshape()`. |
| [Pt Squeeze](pt_squeeze.md) | Removes a dimension at the specified position in the input tensor if it is of size 1. |
+| [Pt Stack](pt_stack.md) | Creates a new axis and stacks two PyTorch tensors along the new axis. |
| [Pt Unsqueeze](pt_unsqueeze.md) | Adds a singleton dimension at the specified position in the input tensor. |
| [Pt View](pt_view.md) | Reshapes a PyTorch tensor into a specified shape using `torch.view()`. |
diff --git a/docs/reference/pt_concat.md b/docs/reference/pt_concat.md
new file mode 100644
index 0000000..2857239
--- /dev/null
+++ b/docs/reference/pt_concat.md
@@ -0,0 +1,19 @@
+# Pt Concat
+Concatenates two PyTorch tensors.
+
+## Input
+| Name | Data type |
+|---|---|
+| tens_a | Tensor |
+| tens_b | Tensor |
+| dim | Int |
+
+## Output
+| Data type |
+|---|
+| Tensor |
+
+
+Category: PyTorch wrapper - Transform
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_stack.md b/docs/reference/pt_stack.md
new file mode 100644
index 0000000..29840a6
--- /dev/null
+++ b/docs/reference/pt_stack.md
@@ -0,0 +1,19 @@
+# Pt Stack
+Creates a new axis and stacks two PyTorch tensors along the new axis.
+
+## Input
+| Name | Data type |
+|---|---|
+| tens_a | Tensor |
+| tens_b | Tensor |
+| dim | Int |
+
+## Output
+| Data type |
+|---|
+| Tensor |
+
+
+Category: PyTorch wrapper - Transform
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/modules/pytorch_wrapper/pt_concat.py b/modules/pytorch_wrapper/pt_concat.py
new file mode 100644
index 0000000..ef0b2f8
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_concat.py
@@ -0,0 +1,42 @@
+from typing import Any, Dict
+import torch
+
+class PtConcat:
+ """
+ Concatenates two PyTorch tensors.
+
+ category: PyTorch wrapper - Transform
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls) -> Dict[str, Any]:
+ """
+ Defines the required input types.
+
+ Returns:
+ Dict[str, Any]: A dictionary of required input types.
+ """
+ return {
+ "required": {
+ "tens_a": ("TENSOR", {}),
+ "tens_b": ("TENSOR", {}),
+ "dim": ("INT", {"default":0, "min":-10, "max":10})
+ }
+ }
+
+ RETURN_TYPES: tuple = ("TENSOR",)
+ FUNCTION: str = "f"
+ CATEGORY: str = "Data Analysis"
+
+ def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor, dim: int) -> tuple:
+ """
+ Concatenates two PyTorch tensors.
+
+ Args:
+ tens_a (torch.Tensor): A PyTorch tensor of any shape.
+ tens_b (torch.Tensor): A PyTorch tensor of shape that supports concatenation with `tens_a`.
+ dim (int): An axis to concatenate on.
+ Returns:
+ tuple: A tuple containing the resultant tensor.
+ """
+ return (torch.concat([tens_a, tens_b], dim=dim),)
diff --git a/modules/pytorch_wrapper/pt_stack.py b/modules/pytorch_wrapper/pt_stack.py
new file mode 100644
index 0000000..b63e5d1
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_stack.py
@@ -0,0 +1,42 @@
+from typing import Any, Dict
+import torch
+
+class PtStack:
+ """
+ Creates a new axis and stacks two PyTorch tensors along the new axis.
+
+ category: PyTorch wrapper - Transform
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls) -> Dict[str, Any]:
+ """
+ Defines the required input types.
+
+ Returns:
+ Dict[str, Any]: A dictionary of required input types.
+ """
+ return {
+ "required": {
+ "tens_a": ("TENSOR", {}),
+ "tens_b": ("TENSOR", {}),
+ "dim": ("INT", {"default":1, "min":-10, "max":10})
+ }
+ }
+
+ RETURN_TYPES: tuple = ("TENSOR",)
+ FUNCTION: str = "f"
+ CATEGORY: str = "Data Analysis"
+
+ def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor, dim: int) -> tuple:
+ """
+ Stacks two PyTorch tensors and create a new axis.
+
+ Args:
+ tens_a (torch.Tensor): A PyTorch tensor of any shape.
+ tens_b (torch.Tensor): A PyTorch tensor of shape that supports concatenation with `tens_a`.
+ dim (int): An axis to stack on.
+ Returns:
+ tuple: A tuple containing the resultant tensor.
+ """
+ return (torch.stack([tens_a, tens_b], dim=dim),)
diff --git a/modules/pytorch_wrapper/ptd_bernoulli.py b/modules/pytorch_wrapper/ptd_bernoulli.py
index beb47c0..d52f465 100644
--- a/modules/pytorch_wrapper/ptd_bernoulli.py
+++ b/modules/pytorch_wrapper/ptd_bernoulli.py
@@ -2,6 +2,30 @@ from typing import Any, Dict
import ast
import torch
from torch.distributions import Bernoulli
+import scipy.stats as scst
+
+
+class BernoulliEx(Bernoulli):
+ """
+ A class to extend Bernoulli to add missing functionality.
+
+ pragma: skip_doc
+ """
+ def cdf(self, x: torch.Tensor):
+ """
+ Computes the cumulative distribution function (CDF) of the Bernoulli distribution.
+
+ Args:
+ x (torch.Tensor): Input tensor containing values where the CDF is evaluated.
+
+ Returns:
+ torch.Tensor: Tensor containing the cumulative probabilities.
+ """
+ a = x.detach().cpu().numpy()
+ probs = self.probs.detach().cpu().numpy()
+ outputs = scst.bernoulli.cdf(a, probs)
+ return torch.tensor(outputs, dtype=torch.float32)
+
class PtdBernoulli:
"""
@@ -52,8 +76,8 @@ class PtdBernoulli:
raise ValueError("You have to specify either probabilities or logits")
if p is not None:
- dist = Bernoulli(probs=torch.tensor(p, dtype=torch.float32))
+ dist = BernoulliEx(probs=torch.tensor(p, dtype=torch.float32))
elif l is not None:
- dist = Bernoulli(logits=torch.tensor(l, dtype=torch.float32))
+ dist = BernoulliEx(logits=torch.tensor(l, dtype=torch.float32))
return (dist,)
\ No newline at end of file
diff --git a/modules/pytorch_wrapper/ptd_beta.py b/modules/pytorch_wrapper/ptd_beta.py
index 8fe4eaa..d25d057 100644
--- a/modules/pytorch_wrapper/ptd_beta.py
+++ b/modules/pytorch_wrapper/ptd_beta.py
@@ -2,6 +2,30 @@ from typing import Any, Dict
import ast
import torch
from torch.distributions import Beta
+import scipy.stats as scst
+
+class BetaEx(Beta):
+ """
+ A class to extend Beta to add missing functionality.
+
+ pragma: skip_doc
+ """
+ def cdf(self, x: torch.Tensor):
+ """
+ Computes the cumulative distribution function (CDF) of the Beta distribution.
+
+ Args:
+ x (torch.Tensor): Input tensor containing values where the CDF is evaluated.
+
+ Returns:
+ torch.Tensor: Tensor containing the cumulative probabilities.
+ """
+ x2 = x.detach().cpu().numpy()
+ a = self.concentration1.detach().cpu().item()
+ b = self.concentration0.detach().cpu().item()
+ outputs = scst.beta.cdf(x2, a, b)
+ return torch.tensor(outputs, dtype=torch.float32).to(x.device)
+
class PtdBeta:
"""
@@ -47,6 +71,6 @@ class PtdBeta:
c1 = ast.literal_eval(alpha)
c2 = ast.literal_eval(beta)
- dist = Beta(
+ dist = BetaEx(
concentration1=torch.tensor(c1, dtype=torch.float32), concentration0=torch.tensor(c2, dtype=torch.float32))
return (dist,)
diff --git a/modules/pytorch_wrapper/ptd_binomial.py b/modules/pytorch_wrapper/ptd_binomial.py
index 0f679b0..171743f 100644
--- a/modules/pytorch_wrapper/ptd_binomial.py
+++ b/modules/pytorch_wrapper/ptd_binomial.py
@@ -2,6 +2,30 @@ from typing import Any, Dict
import ast
import torch
from torch.distributions import Binomial
+import scipy.stats as scst
+
+class BinomialEx(Binomial):
+ """
+ A class to extend Binomial to add missing functionality.
+
+ pragma: skip_doc
+ """
+ def cdf(self, x: torch.Tensor):
+ """
+ Computes the cumulative distribution function (CDF) of the Binomial distribution.
+
+ Args:
+ x (torch.Tensor): Input tensor containing values where the CDF is evaluated.
+
+ Returns:
+ torch.Tensor: Tensor containing the cumulative probabilities.
+ """
+ a = x.detach().cpu().numpy()
+ n = self.total_count.detach().cpu().numpy()
+ probs = self.probs.detach().cpu().numpy()
+ outputs = scst.binom.cdf(a, n, probs)
+ return torch.tensor(outputs, dtype=torch.float32)
+
class PtdBinomial:
"""
@@ -58,11 +82,11 @@ class PtdBinomial:
raise ValueError("You have to specify either probabilities or logits")
if p is not None:
- dist = Binomial(
+ dist = BinomialEx(
total_count=torch.tensor(n, dtype=torch.int64),
probs=torch.tensor(p, dtype=torch.float32))
elif l is not None:
- dist = Binomial(
+ dist = BinomialEx(
total_count=torch.tensor(n, dtype=torch.int64),
logits=torch.tensor(l, dtype=torch.float32))
diff --git a/modules/pytorch_wrapper/ptd_poisson.py b/modules/pytorch_wrapper/ptd_poisson.py
index 4017196..9006f27 100644
--- a/modules/pytorch_wrapper/ptd_poisson.py
+++ b/modules/pytorch_wrapper/ptd_poisson.py
@@ -2,6 +2,29 @@ from typing import Any, Dict
import ast
import torch
from torch.distributions import Poisson
+import scipy.stats as scst
+
+class PoissonEx(Poisson):
+ """
+ A class to extend Poisson to add missing functionality.
+
+ pragma: skip_doc
+ """
+ def cdf(self, x: torch.Tensor):
+ """
+ Computes the cumulative distribution function (CDF) of the Poisson distribution.
+
+ Args:
+ x (torch.Tensor): Input tensor containing values where the CDF is evaluated.
+
+ Returns:
+ torch.Tensor: Tensor containing the cumulative probabilities.
+ """
+ a = x.detach().cpu().numpy()
+ rate = self.rate.detach().cpu().numpy()
+ outputs = scst.poisson.cdf(a, rate)
+ return torch.tensor(outputs, dtype=torch.float32)
+
class PtdPoisson:
"""
@@ -43,6 +66,6 @@ class PtdPoisson:
"""
r = ast.literal_eval(rate)
- dist = Poisson(
+ dist = PoissonEx(
rate=torch.tensor(r, dtype=torch.float32))
return (dist,)
diff --git a/modules/pytorch_wrapper/ptd_student_t.py b/modules/pytorch_wrapper/ptd_student_t.py
index 9accd6d..adefe36 100644
--- a/modules/pytorch_wrapper/ptd_student_t.py
+++ b/modules/pytorch_wrapper/ptd_student_t.py
@@ -2,6 +2,31 @@ from typing import Any, Dict
import ast
import torch
from torch.distributions import StudentT
+import scipy.stats as scst
+
+class StudentTEx(StudentT):
+ """
+ A class to extend StudentT to add missing functionality.
+
+ pragma: skip_doc
+ """
+ def cdf(self, x: torch.Tensor):
+ """
+ Computes the cumulative distribution function (CDF) of the StudentT distribution.
+
+ Args:
+ x (torch.Tensor): Input tensor containing values where the CDF is evaluated.
+
+ Returns:
+ torch.Tensor: Tensor containing the cumulative probabilities.
+ """
+ a = x.detach().cpu().numpy()
+ df = self.df.detach().cpu().item()
+ loc = self.loc.detach().cpu().numpy()
+ scale = self.scale.detach().cpu().numpy()
+ outputs = scst.t.cdf((a - loc)/scale, df)
+ return torch.tensor(outputs, dtype=torch.float32).to(x.device)
+
class PtdStudentT:
"""
@@ -48,13 +73,13 @@ class PtdStudentT:
scale (str): Rate. Standard deviation of the distribution. You can specify a scalar or a tuple of float.
Returns:
- tuple: (torch.distributions.binomial.StudentT)
+ tuple: (torch.distributions.studentT.StudentT)
"""
d = ast.literal_eval(df)
mu = ast.literal_eval(loc)
s = ast.literal_eval(scale)
- dist = StudentT(
+ dist = StudentTEx(
df=torch.tensor(d, dtype=torch.float32),
loc=torch.tensor(mu, dtype=torch.float32),
scale=torch.tensor(s, dtype=torch.float32))
diff --git a/pyproject.toml b/pyproject.toml
index f758b08..2b75f44 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
[project]
name = "pt-wrapper"
description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.1.1"
+version = "3.1.2"
license = {file = "LICENSE"}
[project.urls]
diff --git a/tests/pytorch_wrapper/pt_concat_test.py b/tests/pytorch_wrapper/pt_concat_test.py
new file mode 100644
index 0000000..886a996
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_concat_test.py
@@ -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_concat import PtConcat
+
+
+class TestPtConcat(unittest.TestCase):
+ def setUp(self):
+ """Set up test instance."""
+ self.node = PtConcat()
+
+ # Test cases: (Tensor A, Tensor B, Dim, Expected result)
+ self.test_cases = [
+ (torch.tensor([[1, 2], [3, 4]]), torch.tensor([[5, 6], [7, 8]]), 0, torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])),
+ (torch.tensor([[1, 2], [3, 4]]), torch.tensor([[5, 6], [7, 8]]), 1, torch.tensor([[1, 2, 5, 6], [3, 4, 7, 8]])),
+ (torch.ones(1, 2, 2), torch.zeros(1, 2, 2), 2, torch.tensor([[[1, 1, 0, 0], [1, 1, 0, 0]]])),
+ (torch.ones(1, 1, 2, 2), torch.zeros(1, 1, 2, 2), 3, torch.tensor([[[[1, 1, 0, 0], [1, 1, 0, 0]]]])),
+ ]
+
+ def test_concatenation(self):
+ """Test tensor concatenation along different dimensions."""
+ for tens_a, tens_b, dim, expected in self.test_cases:
+ with self.subTest(tens_a=tens_a, tens_b=tens_b, dim=dim):
+ result, = self.node.f(tens_a, tens_b, dim)
+
+ # Ensure return value is a tensor
+ self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
+
+ # Ensure result matches expected concatenation
+ self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/pytorch_wrapper/pt_stack_test.py b/tests/pytorch_wrapper/pt_stack_test.py
new file mode 100644
index 0000000..2005b26
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_stack_test.py
@@ -0,0 +1,37 @@
+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_stack import PtStack
+
+
+class TestPtStack(unittest.TestCase):
+ def setUp(self):
+ """Set up test instance."""
+ self.node = PtStack()
+
+ # Test cases: (Tensor A, Tensor B, Dim, Expected result)
+ self.test_cases = [
+ (torch.tensor([1, 2]), torch.tensor([5, 6]), 0, torch.tensor([[1, 2], [5, 6]])),
+ (torch.tensor([1, 2]), torch.tensor([5, 6]), 1, torch.tensor([[1, 5], [2, 6]]))
+ ]
+
+ def test_stacking(self):
+ """Test tensor stacking along different dimensions."""
+ for tens_a, tens_b, dim, expected in self.test_cases:
+ with self.subTest(tens_a=tens_a, tens_b=tens_b, dim=dim):
+ result, = self.node.f(tens_a, tens_b, dim)
+
+ # Ensure return value is a tensor
+ self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
+
+ # Ensure result matches expected stacking
+ self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
+
+if __name__ == "__main__":
+ unittest.main()
\ No newline at end of file
diff --git a/tests/pytorch_wrapper/ptdm_cdf_tensor_bernoulli_test.py b/tests/pytorch_wrapper/ptdm_cdf_tensor_bernoulli_test.py
new file mode 100644
index 0000000..5d0957c
--- /dev/null
+++ b/tests/pytorch_wrapper/ptdm_cdf_tensor_bernoulli_test.py
@@ -0,0 +1,67 @@
+import os
+import sys
+import unittest
+import torch
+import scipy.stats as scst
+import numpy as np
+
+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.ptd_bernoulli import BernoulliEx
+from pytorch_wrapper.ptd_bernoulli import PtdBernoulli
+from pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor
+from .utils import assert_equal_for_tensor_and_list_or_scalar
+
+class TestPtdfCdfBernoulli(unittest.TestCase):
+ def setUp(self):
+ """Set up test instance."""
+ self.test_cases = {
+ "Bernoulli": [
+ (0.4, 0, scst.bernoulli.cdf(0, 0.4)),
+ (0.4, 1, scst.bernoulli.cdf(1, 0.4)),
+ (0.4, [0, 1], scst.bernoulli.cdf([0, 1], 0.4)) # 1d array for x
+ ]
+ }
+
+ def test_custom_class_1(self):
+ d = BernoulliEx(probs=0.4)
+ actual = d.cdf(torch.tensor(0.0)).numpy()
+ expected = torch.tensor(0.6).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_2(self):
+ d = BernoulliEx(probs=0.4)
+ actual = d.cdf(torch.tensor(1.0)).numpy()
+ expected = torch.tensor(1.0).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_3(self):
+ d = BernoulliEx(probs=0.4)
+ actual = d.cdf(torch.tensor([0.0, 1.0])).numpy()
+ expected = torch.tensor([0.6, 1.0]).numpy()
+ self.assertTrue(np.allclose(actual, expected, atol=1e-5),
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_distributions(self):
+ """Test probability mass and density functions for the target distribution."""
+ for dist_name, cases in self.test_cases.items(): # For each distribution
+ for case in cases: # For each distribution's test cases
+ with self.subTest(distribution=dist_name, params=case):
+ distribution = PtdBernoulli().f(str(case[0]), "")[0]
+ cdf = PtdmCdfTensor().f(
+ distribution,
+ torch.tensor(case[1]))[0]
+
+ # Ensure return values are tensors
+ self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
+
+ # Check values are equal
+ assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/pytorch_wrapper/ptdm_cdf_tensor_beta.py b/tests/pytorch_wrapper/ptdm_cdf_tensor_beta.py
new file mode 100644
index 0000000..f373dc3
--- /dev/null
+++ b/tests/pytorch_wrapper/ptdm_cdf_tensor_beta.py
@@ -0,0 +1,71 @@
+import os
+import sys
+import unittest
+import torch
+import scipy.stats as scst
+import numpy as np
+
+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.ptd_beta import BetaEx
+from pytorch_wrapper.ptd_beta import PtdBeta
+from pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor
+from .utils import assert_equal_for_tensor_and_list_or_scalar
+
+class TestPtdfCdfBeta(unittest.TestCase):
+ def setUp(self):
+ """Set up test instance."""
+ self.test_cases = {
+ "Beta": [
+ # str alpha, str beta, num x, expected
+ ("2", "8", 0.5, scst.beta.cdf(0.5, 2, 8)),
+ ("2", "8", 0, scst.beta.cdf(0, 2, 8)),
+ ("2", "8", [0, 0.5, 1], scst.beta.cdf([0, 0.5, 1], 2, 8)) # 1d array for x
+ ]
+ }
+
+ def test_custom_class_1(self):
+ d = BetaEx(concentration1=2, concentration0=8)
+ actual = d.cdf(torch.tensor(0.5)).numpy()
+ expected = torch.tensor(0.98046875).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_2(self):
+ d = BetaEx(concentration1=2, concentration0=8)
+ actual = d.cdf(torch.tensor(0.0)).numpy()
+ expected = torch.tensor(0.0).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_3(self):
+ d = BetaEx(concentration1=2, concentration0=8)
+ actual = d.cdf(torch.tensor(1.0)).numpy()
+ expected = torch.tensor(1.0).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+
+ def test_distributions(self):
+ """Test probability mass and density functions for the target distribution."""
+ for dist_name, cases in self.test_cases.items(): # For each distribution
+ for case in cases: # For each distribution's test cases
+ with self.subTest(distribution=dist_name, params=case):
+ distribution = PtdBeta().f(
+ alpha=case[0],
+ beta=case[1])[0]
+ cdf = PtdmCdfTensor().f(
+ distribution,
+ torch.tensor(case[2]))[0]
+
+ # Ensure return values are tensors
+ self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
+
+ # Check values are equal
+ assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/pytorch_wrapper/ptdm_cdf_tensor_binomial_test.py b/tests/pytorch_wrapper/ptdm_cdf_tensor_binomial_test.py
new file mode 100644
index 0000000..0724821
--- /dev/null
+++ b/tests/pytorch_wrapper/ptdm_cdf_tensor_binomial_test.py
@@ -0,0 +1,66 @@
+import os
+import sys
+import unittest
+import torch
+import scipy.stats as scst
+import numpy as np
+
+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.ptd_binomial import BinomialEx
+from pytorch_wrapper.ptd_binomial import PtdBinomial
+from pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor
+from .utils import assert_equal_for_tensor_and_list_or_scalar
+
+class TestPtdfCdfBinomial(unittest.TestCase):
+ def setUp(self):
+ """Set up test instance."""
+ self.test_cases = {
+ "Binomial": [
+ (10.0, 3.0, 0.4, scst.binom.cdf(3.0, 10, 0.4)),
+ (10.0, [3.0, 4.0], 0.4, scst.binom.cdf([3.0, 4.0], 10, 0.4)) # 1d array for x
+ ]
+ }
+
+ def test_custom_class_1(self):
+ d = BinomialEx(total_count=10, probs=0.4)
+ actual = d.cdf(torch.tensor(3.0)).numpy()
+ expected = torch.tensor(0.3822806015999999).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_2(self):
+ d = BinomialEx(total_count=10, probs=0.4)
+ actual = d.cdf(torch.tensor(1.0)).numpy()
+ expected = torch.tensor(0.04635740159999998).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_3(self):
+ d = BinomialEx(total_count=10, probs=0.4)
+ actual = d.cdf(torch.tensor([1.0, 3.0])).numpy()
+ expected = torch.tensor([0.04635740159999998, 0.3822806015999999]).numpy()
+ self.assertTrue(np.allclose(actual, expected, atol=1e-5),
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_distributions(self):
+ """Test probability mass and density functions for the target distribution."""
+ for dist_name, cases in self.test_cases.items(): # For each distribution
+ for case in cases: # For each distribution's test cases
+ with self.subTest(distribution=dist_name, params=case):
+ distribution = PtdBinomial().f(str(case[0]), str(case[2]), "")[0] # n, p
+ cdf = PtdmCdfTensor().f(
+ distribution,
+ torch.tensor(case[1]))[0]
+
+ # Ensure return values are tensors
+ self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
+
+ # Check values are equal
+ assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/pytorch_wrapper/ptdm_cdf_tensor_poisson_test.py b/tests/pytorch_wrapper/ptdm_cdf_tensor_poisson_test.py
new file mode 100644
index 0000000..ecda8e1
--- /dev/null
+++ b/tests/pytorch_wrapper/ptdm_cdf_tensor_poisson_test.py
@@ -0,0 +1,68 @@
+import os
+import sys
+import unittest
+import torch
+import scipy.stats as scst
+import numpy as np
+
+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.ptd_poisson import PoissonEx
+from pytorch_wrapper.ptd_poisson import PtdPoisson
+from pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor
+from .utils import assert_equal_for_tensor_and_list_or_scalar
+
+class TestPtdfCdfPoisson(unittest.TestCase):
+ def setUp(self):
+ """Set up test instance."""
+ self.test_cases = {
+ "Poisson": [
+ # rate, x, expected
+ (3, 2, scst.poisson.cdf(2, 3)),
+ (3, 5, scst.poisson.cdf(5, 3)),
+ (3, [2, 5], scst.poisson.cdf([2, 5], 3)) # 1d array for x
+ ]
+ }
+
+ def test_custom_class_1(self):
+ d = PoissonEx(rate=3)
+ actual = d.cdf(torch.tensor(2)).numpy()
+ expected = torch.tensor(0.42319008112684364).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_2(self):
+ d = PoissonEx(rate=3)
+ actual = d.cdf(torch.tensor(5.0)).numpy()
+ expected = torch.tensor(0.9160820579686966).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_3(self):
+ d = PoissonEx(rate=3)
+ actual = d.cdf(torch.tensor([2.0, 5.0])).numpy()
+ expected = torch.tensor([0.42319008112684364, 0.9160820579686966]).numpy()
+ self.assertTrue(np.allclose(actual, expected, atol=1e-5),
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_distributions(self):
+ """Test probability mass and density functions for the target distribution."""
+ for dist_name, cases in self.test_cases.items(): # For each distribution
+ for case in cases: # For each distribution's test cases
+ with self.subTest(distribution=dist_name, params=case):
+ distribution = PtdPoisson().f(str(case[0]))[0]
+ cdf = PtdmCdfTensor().f(
+ distribution,
+ torch.tensor(case[1]))[0]
+
+ # Ensure return values are tensors
+ self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
+
+ # Check values are equal
+ assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/pytorch_wrapper/ptdm_cdf_tensor_student_t_test.py b/tests/pytorch_wrapper/ptdm_cdf_tensor_student_t_test.py
new file mode 100644
index 0000000..aeb825c
--- /dev/null
+++ b/tests/pytorch_wrapper/ptdm_cdf_tensor_student_t_test.py
@@ -0,0 +1,71 @@
+import os
+import sys
+import unittest
+import torch
+import scipy.stats as scst
+import numpy as np
+
+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.ptd_student_t import StudentTEx
+from pytorch_wrapper.ptd_student_t import PtdStudentT
+from pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor
+from .utils import assert_equal_for_tensor_and_list_or_scalar
+
+class TestPtdfCdfStudentT(unittest.TestCase):
+ def setUp(self):
+ """Set up test instance."""
+ self.test_cases = {
+ "StudentT": [
+ # str df, str loc, str scale, num x, expected
+ ("40.0", "0", "1", 0.5, scst.t.cdf(0.5, 40)),
+ ("40.0", "170", "20", 185, scst.t.cdf((185-170)/20, 40)), # adjust loc and scale
+ ("40.0", "170", "20", [160.0, 185], scst.t.cdf([(160-170)/20, (185-170)/20], 40)) # 1d array for x
+ ]
+ }
+
+ def test_custom_class_1(self):
+ d = StudentTEx(df=40, loc=0, scale=1)
+ actual = d.cdf(torch.tensor(0.5)).numpy()
+ expected = torch.tensor(0.690092632383276).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_2(self):
+ d = StudentTEx(df=40, loc=0, scale=1)
+ actual = d.cdf(torch.tensor(0.9)).numpy()
+ expected = torch.tensor(0.813246561322997).numpy()
+ self.assertAlmostEqual(actual, expected, places=4,
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_custom_class_3(self):
+ d = StudentTEx(df=40, loc=0, scale=1)
+ actual = d.cdf(torch.tensor([0.5, 0.9])).numpy()
+ expected = torch.tensor([0.690092632383276, 0.813246561322997]).numpy()
+ self.assertTrue(np.allclose(actual, expected, atol=1e-5),
+ msg=f"Expected {expected}, got {actual}")
+
+ def test_distributions(self):
+ """Test probability mass and density functions for the target distribution."""
+ for dist_name, cases in self.test_cases.items(): # For each distribution
+ for case in cases: # For each distribution's test cases
+ with self.subTest(distribution=dist_name, params=case):
+ distribution = PtdStudentT().f(
+ df=case[0],
+ loc=case[1],
+ scale=case[2])[0]
+ cdf = PtdmCdfTensor().f(
+ distribution,
+ torch.tensor(case[3]))[0]
+
+ # Ensure return values are tensors
+ self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
+
+ # Check values are equal
+ assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/pytorch_wrapper/ptdm_cdf_test.py b/tests/pytorch_wrapper/ptdm_cdf_test.py
index 7ae9d43..8e9f2df 100644
--- a/tests/pytorch_wrapper/ptdm_cdf_test.py
+++ b/tests/pytorch_wrapper/ptdm_cdf_test.py
@@ -8,11 +8,9 @@ import numpy as np
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.ptd_binomial import PtdBinomial
from pytorch_wrapper.ptdm_cdf import PtdmCdf
from torch.distributions import (
- Binomial, Categorical, Poisson, NegativeBinomial, Multinomial,
+ Categorical, Poisson, NegativeBinomial, Multinomial,
Bernoulli, Beta, Gamma, Normal, Laplace, Exponential, LogNormal,
Pareto, Weibull, Dirichlet, VonMises, StudentT, Uniform, Chi2
)
@@ -20,17 +18,23 @@ from torch.distributions import (
class TestPtdfCdf(unittest.TestCase):
def setUp(self):
- """Set up test instance."""
+ """Set up test instance.
+
+ Following distributions are tested in separate unit test files:
+ * Bernoulli
+ * Binomial
+ * Poisson
+ * StudentT
+ * Beta
+ """
self.test_cases = {
- # "Beta": [(2.0, 5.0, 0.3, scst.beta.cdf(0.3, 2, 5))],
- "Gamma": [(2.0, 1.0, 3.0, scst.gamma.cdf(3.0, 2, scale=1))],
+ "Uniform": [(0.0, 1.0, 0.5, scst.uniform.cdf(0.5, loc=0, scale=1))],
"Normal": [(0.0, 1.0, -1.5, scst.norm.cdf(-1.5, 0, 1))],
+ "Chi2": [(1.0, 2.0, scst.chi2.cdf(2.0, 1))], # df=1, value=2.0
"Exponential": [(1.0, 3.0, scst.expon.cdf(3.0, scale=1))],
+ "Gamma": [(2.0, 1.0, 3.0, scst.gamma.cdf(3.0, 2, scale=1))],
# "Dirichlet": [(torch.tensor([1.0, 1.0, 1.0]), torch.tensor([0.2, 0.5, 0.3]), scst.dirichlet.cdf([0.2, 0.5, 0.3], [1.0, 1.0, 1.0]))],
# "VonMises": [(0.0, 1.0, 1.0, scst.vonmises.cdf(1.0, kappa=1.0, loc=0))],
- # "StudentT": [(1.0, 0.0, scst.t.cdf(0.0, 1))],
- "Uniform": [(0.0, 1.0, 0.5, scst.uniform.cdf(0.5, loc=0, scale=1))],
- "Chi2": [(1.0, 2.0, scst.chi2.cdf(2.0, 1))], # df=1, value=2.0
}
def test_distributions(self):
@@ -38,70 +42,54 @@ class TestPtdfCdf(unittest.TestCase):
for dist_name, cases in self.test_cases.items(): # For each distribution
for case in cases: # For each distribution's test cases
with self.subTest(distribution=dist_name, params=case):
- if dist_name == "Binomial":
- distribution = PtdBinomial().f(str(case[0]), str(case[2]), "")[0] # n, p
- cdf = PtdmCdf().f(distribution, str(case[1]))[0] # k
+ # Instantiate PyTorch distribution directly if not in wrapper
+ if dist_name == "Categorical":
+ distribution = Categorical(case[0])
+ k_tensor = torch.tensor(case[1], dtype=torch.int64)
+ elif dist_name == "NegativeBinomial":
+ distribution = NegativeBinomial(torch.tensor(case[0]), torch.tensor(case[2]))
+ k_tensor = torch.tensor(case[1], dtype=torch.int64)
+ elif dist_name == "Gamma":
+ distribution = Gamma(torch.tensor(case[0]), torch.tensor(case[1]))
+ k_tensor = torch.tensor(case[2], dtype=torch.float32)
+ elif dist_name == "Normal":
+ distribution = Normal(torch.tensor(case[0]), torch.tensor(case[1]))
+ k_tensor = torch.tensor(case[2], dtype=torch.float32)
+ elif dist_name == "Laplace":
+ distribution = Laplace(torch.tensor(case[0]), torch.tensor(case[1]))
+ k_tensor = torch.tensor(case[2], dtype=torch.float32)
+ elif dist_name == "Exponential":
+ distribution = Exponential(torch.tensor(case[0]))
+ k_tensor = torch.tensor(case[1], dtype=torch.float32)
+ elif dist_name == "LogNormal":
+ distribution = LogNormal(torch.tensor(case[0]), torch.tensor(case[1]))
+ k_tensor = torch.tensor(case[2], dtype=torch.float32)
+ elif dist_name == "Pareto":
+ distribution = Pareto(torch.tensor(case[0]), torch.tensor(case[1]))
+ k_tensor = torch.tensor(case[2], dtype=torch.float32)
+ elif dist_name == "Weibull":
+ distribution = Weibull(torch.tensor(1.0), torch.tensor(case[0]))
+ k_tensor = torch.tensor(case[1], dtype=torch.float32)
+ elif dist_name == "Dirichlet":
+ distribution = Dirichlet(case[0])
+ k_tensor = case[1] # Dirichlet expects tensor input
+ elif dist_name == "VonMises":
+ distribution = VonMises(torch.tensor(case[0]), torch.tensor(case[1]))
+ k_tensor = torch.tensor(case[2], dtype=torch.float32)
+ elif dist_name == "Uniform":
+ distribution = Uniform(torch.tensor(case[0]), torch.tensor(case[1]))
+ k_tensor = torch.tensor(case[2], dtype=torch.float32)
+ elif dist_name == "Chi2":
+ distribution = Chi2(torch.tensor(case[0])) # df
+ k_tensor = torch.tensor(case[1], dtype=torch.float32)
else:
- # Instantiate PyTorch distribution directly if not in wrapper
- if dist_name == "Poisson":
- distribution = Poisson(torch.tensor(case[0]))
- k_tensor = torch.tensor(case[1], dtype=torch.int64)
- elif dist_name == "Categorical":
- distribution = Categorical(case[0])
- k_tensor = torch.tensor(case[1], dtype=torch.int64)
- elif dist_name == "NegativeBinomial":
- distribution = NegativeBinomial(torch.tensor(case[0]), torch.tensor(case[2]))
- k_tensor = torch.tensor(case[1], dtype=torch.int64)
- elif dist_name == "Bernoulli":
- distribution = Bernoulli(torch.tensor(case[0]))
- k_tensor = torch.tensor(case[1], dtype=torch.float32)
- elif dist_name == "Beta":
- distribution = Beta(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "Gamma":
- distribution = Gamma(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "Normal":
- distribution = Normal(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "Laplace":
- distribution = Laplace(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "Exponential":
- distribution = Exponential(torch.tensor(case[0]))
- k_tensor = torch.tensor(case[1], dtype=torch.float32)
- elif dist_name == "LogNormal":
- distribution = LogNormal(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "Pareto":
- distribution = Pareto(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "Weibull":
- distribution = Weibull(torch.tensor(1.0), torch.tensor(case[0]))
- k_tensor = torch.tensor(case[1], dtype=torch.float32)
- elif dist_name == "Dirichlet":
- distribution = Dirichlet(case[0])
- k_tensor = case[1] # Dirichlet expects tensor input
- elif dist_name == "VonMises":
- distribution = VonMises(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "StudentT":
- distribution = StudentT(torch.tensor(case[0]))
- k_tensor = torch.tensor(case[1], dtype=torch.float32)
- elif dist_name == "Uniform":
- distribution = Uniform(torch.tensor(case[0]), torch.tensor(case[1]))
- k_tensor = torch.tensor(case[2], dtype=torch.float32)
- elif dist_name == "Chi2":
- distribution = Chi2(torch.tensor(case[0])) # df
- k_tensor = torch.tensor(case[1], dtype=torch.float32)
- else:
- self.fail(f"Unhandled distribution: {dist_name}")
+ self.fail(f"Unhandled distribution: {dist_name}")
- # Call wrapper log-prob method
- if dist_name == "Dirichlet":
- cdf = PtdmCdf().f(distribution, str(k_tensor.tolist()))[0]
- else:
- cdf = PtdmCdf().f(distribution, str(k_tensor.item()))[0]
+ # Call wrapper log-prob method
+ if dist_name == "Dirichlet":
+ cdf = PtdmCdf().f(distribution, str(k_tensor.tolist()))[0]
+ else:
+ cdf = PtdmCdf().f(distribution, str(k_tensor.item()))[0]
# Ensure return values are tensors
self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
diff --git a/tests/pytorch_wrapper/utils.py b/tests/pytorch_wrapper/utils.py
new file mode 100644
index 0000000..0d4d1e0
--- /dev/null
+++ b/tests/pytorch_wrapper/utils.py
@@ -0,0 +1,53 @@
+"""
+Util functions for unit tests.
+"""
+import torch
+import numpy as np
+from typing import Union, List
+import unittest
+
+def assert_equal_for_tensor_and_list_or_scalar(
+ test_inst: unittest.TestCase,
+ actual: torch.Tensor,
+ expected: Union[List[float], float, np.ndarray],
+ message: str = ""
+) -> None:
+ """
+ Asserts that a given tensor's value is approximately equal to the expected value.
+
+ This function supports comparison between a PyTorch tensor and:
+ - A scalar float
+ - A list of floats
+ - A NumPy array
+
+ Args:
+ test_inst (unittest.TestCase): The test instance to perform assertions.
+ actual (torch.Tensor): The tensor containing computed values.
+ expected (Union[List[float], float, np.ndarray]): The expected values.
+ message (str): Additional message to include in assertion errors.
+
+ Raises:
+ ValueError: If `expected` is not a float, list of floats, or a NumPy array.
+ """
+ if isinstance(expected, float):
+ # Convert tensor to scalar float
+ actual_value = actual.item()
+ test_inst.assertAlmostEqual(
+ actual_value, expected, places=4,
+ msg=f"{message}: Expected {expected}, got {actual_value}"
+ )
+
+ elif isinstance(expected, list) or isinstance(expected, np.ndarray):
+ # Convert tensor to numpy array for comparison
+ actual_array = actual.cpu().numpy()
+ expected_array = np.array(expected, dtype=actual_array.dtype)
+
+ test_inst.assertTrue(
+ np.allclose(actual_array, expected_array, atol=1e-5),
+ msg=f"{message}: Expected {expected_array}, got {actual_array}"
+ )
+
+ else:
+ raise ValueError(
+ f"Unsupported data type for `expected`: {type(expected)}"
+ )