From 6e1a99cce32f5545b311caab67fb235d53513fb7 Mon Sep 17 00:00:00 2001
From: Aarsh-Wankar <23110003@iitgn.ac.in>
Date: Sat, 22 Mar 2025 00:15:46 +0530
Subject: [PATCH] Add kn function as helper for modified Bessel function of the
 second kind and corresponding tests

---
 pytensor/tensor/math.py         |  6 ++++++
 tests/tensor/test_math_scipy.py | 16 +++++++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index 2ba2ab1587..714f597b32 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -2453,6 +2453,11 @@ def kv(v, x):
     return kve(v, x) * exp(-x)
 
 
+def kn(n, x):
+    """Modified Bessel function of the second kind of integer order v."""
+    return kv(n, x)
+
+
 @scalar_elemwise
 def sigmoid(x):
     """Logistic sigmoid function (1 / (1 + exp(-x)), also known as expit or inverse logit"""
@@ -4337,6 +4342,7 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     "i1",
     "iv",
     "ive",
+    "kn",
     "kv",
     "kve",
     "sigmoid",
diff --git a/tests/tensor/test_math_scipy.py b/tests/tensor/test_math_scipy.py
index 8f70950206..e15293e979 100644
--- a/tests/tensor/test_math_scipy.py
+++ b/tests/tensor/test_math_scipy.py
@@ -18,7 +18,7 @@
 from pytensor import tensor as pt
 from pytensor.compile.mode import get_default_mode
 from pytensor.configdefaults import config
-from pytensor.tensor import gammaincc, inplace, kv, kve, vector
+from pytensor.tensor import gammaincc, inplace, kn, kv, kve, vector
 from tests import unittest_tools as utt
 from tests.tensor.utils import (
     _good_broadcast_unary_chi2sf,
@@ -1220,3 +1220,17 @@ def test_kv():
         out.eval({v: test_v, x: test_x}),
         scipy.special.kv(test_v[:, None], test_x[None, :]),
     )
+
+
+def test_kn():
+    n = vector("n")
+    x = vector("x")
+
+    out = kn(n[:, None], x[None, :])
+    test_n = np.array([-3, 4, 0, 5], dtype=n.type.dtype)
+    test_x = np.linspace(0, 512, 10, dtype=x.type.dtype)
+
+    np.testing.assert_allclose(
+        out.eval({n: test_n, x: test_x}),
+        scipy.special.kn(test_n[:, None], test_x[None, :]),
+    )