Added cdf for some nodes. Added stack and concat nodes.

__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",

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()`. |

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 |
+
+<HR>
+Category: PyTorch wrapper - Transform
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

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 |
+
+<HR>
+Category: PyTorch wrapper - Transform
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

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),)

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),)

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,)

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,)

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

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,)

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

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]

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()

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()

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()

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()

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()

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()

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()

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

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)}"
+        )