Added distribution method nodes that take 2d tensor as input

__init__.py

@@ -124,9 +124,13 @@ from .modules.pytorch_wrapper.ptd_poisson import PtdPoisson
 from .modules.pytorch_wrapper.ptd_student_t import PtdStudentT
 from .modules.pytorch_wrapper.ptd_uniform import PtdUniform
 from .modules.pytorch_wrapper.ptdm_cdf import PtdmCdf
+from .modules.pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor
 from .modules.pytorch_wrapper.ptdm_log_prob import PtdmLogProb
+from .modules.pytorch_wrapper.ptdm_log_prob_tensor import PtdmLogProbTensor
 from .modules.pytorch_wrapper.ptdm_pdf import PtdmPdf
+from .modules.pytorch_wrapper.ptdm_pdf_tensor import PtdmPdfTensor
 from .modules.pytorch_wrapper.ptdm_pmf import PtdmPmf
+from .modules.pytorch_wrapper.ptdm_pmf_tensor import PtdmPmfTensor
 from .modules.pytorch_wrapper.ptdm_sample import PtdmSample
 from .modules.pytorch_wrapper.ptf_gelu import PtfGELU
 from .modules.pytorch_wrapper.ptf_leaky_relu import PtfLeakyReLU
@@ -294,9 +298,13 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtdStudentT": PtdStudentT,
     "PtdUniform": PtdUniform,
     "PtdmCdf": PtdmCdf,
+    "PtdmCdfTensor": PtdmCdfTensor,
     "PtdmLogProb": PtdmLogProb,
+    "PtdmLogProbTensor": PtdmLogProbTensor,
     "PtdmPdf": PtdmPdf,
+    "PtdmPdfTensor": PtdmPdfTensor,
     "PtdmPmf": PtdmPmf,
+    "PtdmPmfTensor": PtdmPmfTensor,
     "PtdmSample": PtdmSample,
     "PtfGELU": PtfGELU,
     "PtfLeakyReLU": PtfLeakyReLU,
@@ -464,9 +472,13 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtdStudentT": "Ptd Student T",
     "PtdUniform": "Ptd Uniform",
     "PtdmCdf": "Ptdm Cdf",
+    "PtdmCdfTensor": "Ptdm Cdf Tensor",
     "PtdmLogProb": "Ptdm Log Prob",
+    "PtdmLogProbTensor": "Ptdm Log Prob Tensor",
     "PtdmPdf": "Ptdm Pdf",
+    "PtdmPdfTensor": "Ptdm Pdf Tensor",
     "PtdmPmf": "Ptdm Pmf",
+    "PtdmPmfTensor": "Ptdm Pmf Tensor",
     "PtdmSample": "Ptdm Sample",
     "PtfGELU": "Ptf GELU",
     "PtfLeakyReLU": "Ptf Leaky ReLU",

docs/reference/node_reference.md

@@ -75,9 +75,13 @@
 | [Ptd Student T](ptd_student_t.md) | Instantiates a StudentT distribution object. |
 | [Ptd Uniform](ptd_uniform.md) | Instantiates a Uniform distribution object. |
 | [Ptdm Cdf](ptdm_cdf.md) | Computes the cumulative distribution function for the input distribution. |
+| [Ptdm Cdf Tensor](ptdm_cdf_tensor.md) | Computes the cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. |
 | [Ptdm Log Prob](ptdm_log_prob.md) | Computes the log of probability for the input distribution. |
+| [Ptdm Log Prob Tensor](ptdm_log_prob_tensor.md) | Computes the log of probability for the input distribution.  This nodes accepts a tensor so it can be used to compute log of probability for multiple values contained in a tensor. |
 | [Ptdm Pdf](ptdm_pdf.md) | Computes the probability density for the input distribution. |
+| [Ptdm Pdf Tensor](ptdm_pdf_tensor.md) | Computes the probability density for the input distribution.  This nodes accepts a tensor so it can be used to compute pdf for multiple values contained in a tensor. |
 | [Ptdm Pmf](ptdm_pmf.md) | Computes the probability for the input distribution. |
+| [Ptdm Pmf Tensor](ptdm_pmf_tensor.md) | Computes the probability for the input distribution.  This nodes accepts a tensor so it can be used to compute pmf for multiple values contained in a tensor. |
 | [Ptdm Sample](ptdm_sample.md) | Samples from the input distribution. |
 ## PyTorch wrapper - Image processing
 | Node | Description |

docs/reference/ptdm_cdf.md

@@ -6,13 +6,13 @@ Cdf is not supported for all distributions in PyTorch.
 
 Args:
         distribution (torch.distributions.distribution.Distribution): Distribution.
-        v (float): Value
+        x (float): Value
 
 ## Input
 | Name | Data type |
 |---|---|
 | distribution | Ptdistribution |
-| v | String |
+| x | String |
 
 ## Output
 | Data type |

docs/reference/ptdm_cdf_tensor.md

@@ -0,0 +1,25 @@
+# Ptdm Cdf Tensor
+Computes the cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor.
+
+**Note**  
+Cdf is not supported for all distributions in PyTorch.
+
+Args:
+        distribution (torch.distributions.distribution.Distribution): Distribution.
+        tens (torch.Tensor): Value(s) in Tensor.
+
+## Input
+| Name | Data type |
+|---|---|
+| distribution | Ptdistribution |
+| tens | Tensor |
+
+## Output
+| Data type |
+|---|
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Distribution
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptdm_log_prob.md

@@ -9,7 +9,7 @@ Args:
 | Name | Data type |
 |---|---|
 | distribution | Ptdistribution |
-| value | String |
+| x | String |
 
 ## Output
 | Data type |

docs/reference/ptdm_log_prob_tensor.md

@@ -0,0 +1,22 @@
+# Ptdm Log Prob Tensor
+Computes the log of probability for the input distribution.  This nodes accepts a tensor so it can be used to compute log of probability for multiple values contained in a tensor.
+
+Args:
+        distribution (torch.distributions.distribution.Distribution): Distribution.  
+        tens (torch.Tensor): Value(s) in Tensor.
+
+## Input
+| Name | Data type |
+|---|---|
+| distribution | Ptdistribution |
+| tens | Tensor |
+
+## Output
+| Data type |
+|---|
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Distribution
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptdm_pdf.md

@@ -3,13 +3,13 @@ Computes the probability density for the input distribution.
 
 Args:
         distribution (torch.distributions.distribution.Distribution): Distribution.
-        v (float): Value
+        x (float): Value
 
 ## Input
 | Name | Data type |
 |---|---|
 | distribution | Ptdistribution |
-| v | String |
+| x | String |
 
 ## Output
 | Data type |

docs/reference/ptdm_pdf_tensor.md

@@ -0,0 +1,22 @@
+# Ptdm Pdf Tensor
+Computes the probability density for the input distribution.  This nodes accepts a tensor so it can be used to compute pdf for multiple values contained in a tensor.
+
+Args:
+        distribution (torch.distributions.distribution.Distribution): Distribution.
+        tens (torch.Tensor): Value(s) in Tensor.
+
+## Input
+| Name | Data type |
+|---|---|
+| distribution | Ptdistribution |
+| tens | Tensor |
+
+## Output
+| Data type |
+|---|
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Distribution
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptdm_pmf_tensor.md

@@ -0,0 +1,22 @@
+# Ptdm Pmf Tensor
+Computes the probability for the input distribution.  This nodes accepts a tensor so it can be used to compute pmf for multiple values contained in a tensor.
+
+Args:
+        distribution (torch.distributions.distribution.Distribution): Distribution.
+        tens (torch.Tensor): k(s) in Tensor.
+
+## Input
+| Name | Data type |
+|---|---|
+| distribution | Ptdistribution |
+| tens | Tensor |
+
+## Output
+| Data type |
+|---|
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Distribution
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

modules/pytorch_wrapper/ptdm_cdf.py

@@ -12,7 +12,7 @@ class PtdmCdf:
 
     Args:
             distribution (torch.distributions.distribution.Distribution): Distribution.
-            v (float): Value
+            x (float): Value
 
     category: PyTorch wrapper - Distribution
     """
@@ -28,7 +28,7 @@ class PtdmCdf:
         return {
             "required": {
                 "distribution": ("PTDISTRIBUTION", {}),
-                "v": ("STRING", {"default": ""})
+                "x": ("STRING", {"default": ""})
             }
         }
 
@@ -36,18 +36,18 @@ class PtdmCdf:
     FUNCTION: str = "f"
     CATEGORY: str = "Distribution"
 
-    def f(self, distribution: torch.distributions.Distribution, v: str) -> tuple:
+    def f(self, distribution: torch.distributions.Distribution, x: str) -> tuple:
         """
         Computes the cumulative distribution function for the input distribution.
 
         Args:
             distribution (torch.distributions.Distribution): Distribution.
-            v (str): Value
+            x (str): Value
 
         Returns:
             tuple containing the probability
         """
-        v2 = ast.literal_eval(v)
+        v2 = ast.literal_eval(x)
         if isinstance(v2, int) is False and isinstance(v2, float) is False :
             if isinstance(distribution, Dirichlet) is False:
                 raise ValueError("You need to specify a float or an int.")

modules/pytorch_wrapper/ptdm_cdf_tensor.py

@@ -0,0 +1,49 @@
+from typing import Any, Dict
+import torch
+
+
+class PtdmCdfTensor:
+    """
+    Computes the cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor.
+
+    **Note**  
+    Cdf is not supported for all distributions in PyTorch.
+
+    Args:
+            distribution (torch.distributions.distribution.Distribution): Distribution.
+            tens (torch.Tensor): Value(s) in Tensor.
+
+    category: PyTorch wrapper - Distribution
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying the required input types.
+        """
+        return {
+            "required": {
+                "distribution": ("PTDISTRIBUTION", {}),
+                "tens": ("TENSOR", {})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("TENSOR",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Distribution"
+
+    def f(self, distribution: torch.distributions.Distribution, tens: torch.Tensor) -> tuple:
+        """
+        Computes the cumulative distribution function for the input distribution.
+
+        Args:
+            distribution (torch.distributions.Distribution): Distribution.
+            tens (torch.Tensor): Value(s) in Tensor.
+
+        Returns:
+            tuple containing the probability
+        """
+        return (distribution.cdf(tens),)

modules/pytorch_wrapper/ptdm_log_prob.py

@@ -25,7 +25,7 @@ class PtdmLogProb:
         return {
             "required": {
                 "distribution": ("PTDISTRIBUTION", {}),
-                "value": ("STRING", {"default": ""})
+                "x": ("STRING", {"default": ""})
             }
         }
 
@@ -35,27 +35,27 @@ class PtdmLogProb:
 
     def f(self,
           distribution:torch.distributions.distribution.Distribution,
-          value: str) -> tuple:
+          x: str) -> tuple:
         """
         Computes the probability for the input distribution.
 
         Args:
             distribution (torch.distributions.distribution.Distribution): Distribution.
-            value (Union[int,float])): Value or number of successes.
+            x (Union[int,float])): Value or number of successes.
 
         Returns:
             tuple containing the log of probability
         """
 
-        k = ast.literal_eval(value)
-        if isinstance(k, int) is False and isinstance(k, float) is False :
+        v = ast.literal_eval(x)
+        if isinstance(v, int) is False and isinstance(v, float) is False :
             if isinstance(distribution, Dirichlet) is False:
                 raise ValueError("You need to specify a float or an int.")
 
         if isinstance(distribution, (Poisson, Categorical, Multinomial, Binomial, NegativeBinomial)):
-            tens = torch.tensor(k, dtype=torch.int64)
+            tens = torch.tensor(v, dtype=torch.int64)
         else:
-            tens = torch.tensor(k, dtype=torch.float32)
+            tens = torch.tensor(v, dtype=torch.float32)
 
         return (distribution.log_prob(tens),)
     

modules/pytorch_wrapper/ptdm_log_prob_tensor.py

@@ -0,0 +1,49 @@
+from typing import Any, Dict
+import torch
+
+
+class PtdmLogProbTensor:
+    """
+    Computes the log of probability for the input distribution.  This nodes accepts a tensor so it can be used to compute log of probability for multiple values contained in a tensor.
+
+    Args:
+            distribution (torch.distributions.distribution.Distribution): Distribution.  
+            tens (torch.Tensor): Value(s) in Tensor.
+
+    category: PyTorch wrapper - Distribution
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying the required input types.
+        """
+        return {
+            "required": {
+                "distribution": ("PTDISTRIBUTION", {}),
+                "tens": ("TENSOR", {})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("TENSOR",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Distribution"
+
+    def f(self,
+          distribution:torch.distributions.distribution.Distribution,
+          tens: torch.Tensor) -> tuple:
+        """
+        Computes the probability for the input distribution.
+
+        Args:
+            distribution (torch.distributions.distribution.Distribution): Distribution.
+            tens (torch.Tensor): Value(s) in Tensor.
+
+        Returns:
+            tuple containing the log of probability
+        """
+        return (distribution.log_prob(tens),)
+    

modules/pytorch_wrapper/ptdm_pdf.py

@@ -9,7 +9,7 @@ class PtdmPdf:
 
     Args:
             distribution (torch.distributions.distribution.Distribution): Distribution.
-            v (float): Value
+            x (float): Value
 
     category: PyTorch wrapper - Distribution
     """
@@ -25,7 +25,7 @@ class PtdmPdf:
         return {
             "required": {
                 "distribution": ("PTDISTRIBUTION", {}),
-                "v": ("STRING", {"default": ""})
+                "x": ("STRING", {"default": ""})
             }
         }
 
@@ -33,18 +33,18 @@ class PtdmPdf:
     FUNCTION: str = "f"
     CATEGORY: str = "Distribution"
 
-    def f(self, distribution: torch.distributions.Distribution, v: str) -> tuple:
+    def f(self, distribution: torch.distributions.Distribution, x: str) -> tuple:
         """
         Computes the probability for the input distribution.
 
         Args:
             distribution (torch.distributions.Distribution): Distribution.
-            v (str): Value
+            x (str): Value
 
         Returns:
             tuple containing the probability
         """
-        v2 = ast.literal_eval(v)
+        v2 = ast.literal_eval(x)
         if isinstance(v2, int) is False and isinstance(v2, float) is False :
             if isinstance(distribution, Dirichlet) is False:
                 raise ValueError("You need to specify a float or an int.")

modules/pytorch_wrapper/ptdm_pdf_tensor.py

@@ -0,0 +1,46 @@
+from typing import Any, Dict
+import torch
+
+
+class PtdmPdfTensor:
+    """
+    Computes the probability density for the input distribution.  This nodes accepts a tensor so it can be used to compute pdf for multiple values contained in a tensor.
+
+    Args:
+            distribution (torch.distributions.distribution.Distribution): Distribution.
+            tens (torch.Tensor): Value(s) in Tensor.
+
+    category: PyTorch wrapper - Distribution
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying the required input types.
+        """
+        return {
+            "required": {
+                "distribution": ("PTDISTRIBUTION", {}),
+                "tens": ("TENSOR", {})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("TENSOR",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Distribution"
+
+    def f(self, distribution: torch.distributions.Distribution, tens: torch.Tensor) -> tuple:
+        """
+        Computes the probability for the input distribution.
+
+        Args:
+            distribution (torch.distributions.Distribution): Distribution.
+            tens (torch.Tensor): Value(s) in Tensor.
+
+        Returns:
+            tuple containing the probability
+        """
+        return (torch.exp(distribution.log_prob(tens)),)

modules/pytorch_wrapper/ptdm_pmf_tensor.py

@@ -0,0 +1,46 @@
+from typing import Any, Dict
+import torch
+
+
+class PtdmPmfTensor:
+    """
+    Computes the probability for the input distribution.  This nodes accepts a tensor so it can be used to compute pmf for multiple values contained in a tensor.
+
+    Args:
+            distribution (torch.distributions.distribution.Distribution): Distribution.
+            tens (torch.Tensor): k(s) in Tensor.
+
+    category: PyTorch wrapper - Distribution
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying the required input types.
+        """
+        return {
+            "required": {
+                "distribution": ("PTDISTRIBUTION", {}),
+                "tens": ("TENSOR", {})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("TENSOR",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Distribution"
+
+    def f(self, distribution: torch.distributions.Distribution, tens: torch.Tensor) -> tuple:
+        """
+        Computes the probability for the input distribution.
+
+        Args:
+            distribution (torch.distributions.Distribution): Distribution.
+            tens (torch.Tensor): Value(s) in Tensor.
+
+        Returns:
+            tuple containing the probability
+        """
+        return (torch.exp(distribution.log_prob(tens)),)

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.0"
+version = "3.1.1"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/ptdm_cdf_tensor_test.py

@@ -0,0 +1,61 @@
+import os
+import sys
+import unittest
+import numpy as np
+import torch
+import scipy.stats as scst
+from torch.distributions import Exponential
+
+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.ptdm_cdf_tensor import PtdmCdfTensor
+
+
+class TestPtdfCdfTensor(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.test_cases = {
+            "Exponential": [
+                (1.0, 3.0, scst.expon.cdf(3.0, scale=1)),
+                (1.0, [3.0, 4.0], scst.expon.cdf(np.array([3.0, 4.0], dtype=np.float32), scale=1))
+            ],
+        }
+
+    def test_distributions(self):
+        """Test CDF."""
+        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 == "Exponential":
+                        distribution = Exponential(torch.tensor(case[0]))
+                        k_tensor = torch.tensor(case[1], dtype=torch.float32)
+                    else:
+                        raise ValueError("Only exponential is supported in this unit test.")
+
+                    cdf = PtdmCdfTensor().f(distribution, k_tensor)[0]
+
+                    # Ensure return values are tensors
+                    self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
+
+                    # Convert tensors to floats for comparison
+                    if isinstance(case[1], float):
+                        cdf = cdf.item()
+
+                        # Check if computed values are approximately equal to expected values
+                        self.assertAlmostEqual(cdf, case[-1], places=4, 
+                                            msg=f"{dist_name}: Expected {case[-1]}, got {cdf}")
+
+                    elif isinstance(case[1], list):
+                        cdf = cdf.numpy()
+                        self.assertTrue(
+                           np.allclose(cdf, case[-1], atol=1e-5)
+                        )
+
+                    else:
+                        raise ValueError("Unsupported data type.")
+                    
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptdm_log_prob_tensor_test.py

@@ -0,0 +1,61 @@
+import os
+import sys
+import unittest
+import numpy as np
+import torch
+import scipy.stats as scst
+from torch.distributions import Exponential
+
+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.ptdm_log_prob_tensor import PtdmLogProbTensor
+
+
+class TestPtdfLogProbTensor(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.test_cases = {
+            "Exponential": [
+                (1.0, 3.0, scst.expon.logpdf(3.0, scale=1)),
+                (1.0, [3.0, 4.0], scst.expon.logpdf(np.array([3.0, 4.0], dtype=np.float32), scale=1))
+            ],
+        }
+
+    def test_distributions(self):
+        """Test log prob."""
+        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 == "Exponential":
+                        distribution = Exponential(torch.tensor(case[0]))
+                        k_tensor = torch.tensor(case[1], dtype=torch.float32)
+                    else:
+                        raise ValueError("Only exponential is supported in this unit test.")
+
+                    log_prob = PtdmLogProbTensor().f(distribution, k_tensor)[0]
+
+                    # Ensure return values are tensors
+                    self.assertTrue(isinstance(log_prob, torch.Tensor), "Probability output is not a tensor")
+
+                    # Convert tensors to floats for comparison
+                    if isinstance(case[1], float):
+                        log_prob = log_prob.item()
+
+                        # Check if computed values are approximately equal to expected values
+                        self.assertAlmostEqual(log_prob, case[-1], places=4, 
+                                            msg=f"{dist_name}: Expected {case[-1]}, got {log_prob}")
+
+                    elif isinstance(case[1], list):
+                        log_prob = log_prob.numpy()
+                        self.assertTrue(
+                           np.allclose(log_prob, case[-1], atol=1e-5)
+                        )
+
+                    else:
+                        raise ValueError("Unsupported data type.")
+                    
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptdm_pdf_tensor_test.py

@@ -0,0 +1,61 @@
+import os
+import sys
+import unittest
+import numpy as np
+import torch
+import scipy.stats as scst
+from torch.distributions import Exponential
+
+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.ptdm_pdf_tensor import PtdmPdfTensor
+
+
+class TestPtdfPdfTensor(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.test_cases = {
+            "Exponential": [
+                (1.0, 3.0, scst.expon.pdf(3.0, scale=1)),
+                (1.0, [3.0, 4.0], scst.expon.pdf(np.array([3.0, 4.0], dtype=np.float32), scale=1))
+            ],
+        }
+
+    def test_distributions(self):
+        """Test pdf."""
+        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 == "Exponential":
+                        distribution = Exponential(torch.tensor(case[0]))
+                        k_tensor = torch.tensor(case[1], dtype=torch.float32)
+                    else:
+                        raise ValueError("Only exponential is supported in this unit test.")
+
+                    pdf = PtdmPdfTensor().f(distribution, k_tensor)[0]
+
+                    # Ensure return values are tensors
+                    self.assertTrue(isinstance(pdf, torch.Tensor), "Probability output is not a tensor")
+
+                    # Convert tensors to floats for comparison
+                    if isinstance(case[1], float):
+                        pdf = pdf.item()
+
+                        # Check if computed values are approximately equal to expected values
+                        self.assertAlmostEqual(pdf, case[-1], places=4, 
+                                            msg=f"{dist_name}: Expected {case[-1]}, got {pdf}")
+
+                    elif isinstance(case[1], list):
+                        pdf = pdf.numpy()
+                        self.assertTrue(
+                           np.allclose(pdf, case[-1], atol=1e-5)
+                        )
+
+                    else:
+                        raise ValueError("Unsupported data type.")
+                    
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptdm_pmf_tensor_test.py

@@ -0,0 +1,61 @@
+import os
+import sys
+import unittest
+import numpy as np
+import torch
+import scipy.stats as scst
+from torch.distributions import Poisson
+
+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.ptdm_pmf_tensor import PtdmPmfTensor
+
+
+class TestPtdfPmfTensor(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.test_cases = {
+            "Poisson": [
+                (5, 3.0, scst.poisson.pmf(3.0, mu=5)),
+                (5, [3.0, 4.0], scst.poisson.pmf(np.array([3.0, 4.0], dtype=np.float32), mu=5))
+            ],
+        }
+
+    def test_distributions(self):
+        """Test pmf."""
+        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 == "Poisson":
+                        distribution = Poisson(torch.tensor(case[0]))
+                        k_tensor = torch.tensor(case[1], dtype=torch.float32)
+                    else:
+                        raise ValueError("Only Poisson is supported in this unit test.")
+
+                    pmf = PtdmPmfTensor().f(distribution, k_tensor)[0]
+
+                    # Ensure return values are tensors
+                    self.assertTrue(isinstance(pmf, torch.Tensor), "Probability output is not a tensor")
+
+                    # Convert tensors to floats for comparison
+                    if isinstance(case[1], float):
+                        pmf = pmf.item()
+
+                        # Check if computed values are approximately equal to expected values
+                        self.assertAlmostEqual(pmf, case[-1], places=4, 
+                                            msg=f"{dist_name}: Expected {case[-1]}, got {pmf}")
+
+                    elif isinstance(case[1], list):
+                        pmf = pmf.numpy()
+                        self.assertTrue(
+                           np.allclose(pmf, case[-1], atol=1e-5)
+                        )
+
+                    else:
+                        raise ValueError("Unsupported data type.")
+                    
+
+if __name__ == "__main__":
+    unittest.main()