Added distribution method nodes that take 2d tensor as input

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